def solve(): N, K = map(int, input().split()) max_heap = OrderHeap([], lambda key: -key) unique = set() cnt = defaultdict(int) max_heap.push(N) cnt[N] = 1 while len(max_heap) > 0: val = max_heap.pop() nr = cnt[val] if K <= nr: return (val - 1 - (val-1)//2, (val - 1) // 2) else: K -= nr unique.add(val) l = [(val-1)//2, val - 1 - (val-1)//2] for el in l: if el: if cnt[el] is 0: max_heap.push(el) cnt[el] += nr
def test_push(self): heap = OrderHeap([], key=self.key) wanted = set() not_wanted = set(ascii_uppercase) for c in reversed(ascii_uppercase): heap.push(c) wanted.add(c) not_wanted.remove(c) self.assertHeap(wanted, not_wanted, heap) self.assertHeap(ascii_uppercase, [], heap)
def test_pushpop(self): heap = OrderHeap(reversed(ascii_lowercase), key=self.key) wanted = set(ascii_lowercase) not_wanted = set() for u, l in reversed(list(zip(ascii_uppercase, ascii_lowercase))): self.assertEqual(l, heap.pushpop(u)) wanted.add(u) wanted.remove(l) not_wanted.add(l) self.assertHeap(wanted, not_wanted, heap) self.assertHeap(ascii_uppercase, ascii_lowercase, heap)
def test_pop(self): heap = OrderHeap(reversed(ascii_uppercase), key=self.key) wanted = set(ascii_uppercase) not_wanted = set() sorted_items = [] for c in reversed(ascii_uppercase): self.assertEqual(c, heap.pop()) wanted.remove(c) not_wanted.add(c) self.assertHeap(wanted, not_wanted, heap) sorted_items.append(c) self.assertSequenceEqual(list(reversed(ascii_uppercase)), sorted_items) self.assertHeap([], ascii_uppercase, heap)
def test_repr(self): heap = OrderHeap(reversed(ascii_uppercase), key=self.key) self.assertHeap(heap, [], eval(re.sub(r'key=<.*>', 'key=self.key', repr(heap))))
def test_remove_not_implemented(self): heap = OrderHeap(reversed(ascii_uppercase), key=self.key) self.assertRaises(NotImplementedError, heap.remove, 'A')
def test_pushpop_on_empty_heap(self): self.assertEqual('A', OrderHeap(key=self.key).pushpop('A'))
def test_peek(self): heap = OrderHeap(reversed(ascii_uppercase), key=self.key) self.assertEqual('Z', heap.peek())
def test_check_variant_invalid(self): heap = OrderHeap(ascii_uppercase, key=self.key) heap[3] = (self.key('t'), 't') self.assertRaises(InvalidHeapError, heap.check)
def test_check(self): OrderHeap([], key=self.key).check() OrderHeap(ascii_uppercase, key=self.key).check() OrderHeap(reversed(ascii_uppercase), key=self.key).check()
def test_init(self): self.assertHeap([], [], OrderHeap(key=self.key)) self.assertHeap([], [], OrderHeap([], key=self.key)) self.assertHeap(ascii_uppercase, [], OrderHeap(ascii_uppercase, key=self.key))