{"id":"dsa-trees-heaps","slug":"heaps","title":"Heaps and heapq","text":"Some programs need the smallest item again and again while new items keep arriving: a printer that prints the shortest job next, a game that runs the next event in time order, an emergency room that treats the most urgent patient first. A collection like that is a **priority queue**. With a plain list, `min` finds the smallest in O(n) every time. A list kept sorted does no better in Big-O terms: each new item costs O(n), whether `bisect.insort` puts it in its place (every item after it moves, which is quick work in C but still O(n)) or you append it and sort again (Timsort sees that the list is almost in order, but still looks at every item). A **heap** adds an item in O(log n), removes the smallest in O(log n), and shows the smallest in O(1).\n\n### A tree in a list\n\nA **binary heap** is a binary tree with two rules:\n\n- **Shape:** every level is full, except maybe the last, which fills from the left. So the height is about log₂ n.\n- **Order:** every parent is smaller than or equal to its children. So the smallest value is at the root. (Unlike in a BST, nothing orders the left side against the right.)\n\nWith no gaps in the tree, it fits in a plain list, level by level, without any `Node` objects: the root at index 0, the children of index `i` at `2 * i + 1` and `2 * i + 2`, and its parent at `(i - 1) // 2`.\n\n```text\nindex:  0  1  2  3  4\n       [1, 3, 8, 5, 4]\n\n           1\n         /   \\\n        3     8\n       / \\\n      5   4\n```\n\nTo **push** a value, append it (the next free place on the last level) and **sift it up**: while it's smaller than its parent, swap the two. To **pop** the smallest, take the root, move the last item into its place and **sift it down**: while it's greater than the smaller of its children, swap it with that child. Each walks one path between the root and the bottom: O(log n).\n\nWhy the smaller child? After the swap, that child is the parent of its old sibling, and it is no greater than that sibling, so the order rule holds there. In the heap above, popping 1 moves 4 to the root, above 3 and 8. Swapping 4 with 3 gives `[3, 4, 8, 5]`, a heap again. Swapping it with 8, the larger child, would give `[8, 3, 4, 5]`, with 8 above 3. Nothing stops with an error, but the next pop returns 8 instead of 3. Always swapping with the left child goes wrong in the same way whenever the right child is the smaller one. Sifting up keeps the rule for a similar reason: the value that moves up is smaller than its old parent, which was no greater than its other child.\n\n```python\ndef push(heap, value):\n    heap.append(value)\n    i = len(heap) - 1\n    while i > 0:\n        parent = (i - 1) // 2\n        if heap[i] >= heap[parent]:\n            break\n        heap[i], heap[parent] = heap[parent], heap[i]\n        i = parent\n\n\ndef pop(heap):\n    smallest = heap[0]\n    last = heap.pop()\n    if heap:\n        heap[0] = last\n        i = 0\n        while 2 * i + 1 < len(heap):\n            child = 2 * i + 1\n            if child + 1 < len(heap) and heap[child + 1] < heap[child]:\n                child += 1\n            if heap[i] <= heap[child]:\n                break\n            heap[i], heap[child] = heap[child], heap[i]\n            i = child\n    return smallest\n\n\nheap = []\nfor value in [5, 3, 8, 1, 4]:\n    push(heap, value)\nprint(heap)\nprint([pop(heap) for _ in range(5)])\n```\n\n```output\n[1, 3, 8, 5, 4]\n[1, 3, 4, 5, 8]\n```\n\nThe list isn't sorted, and it doesn't need to be: a heap only keeps the smallest in front, which is far less work than keeping everything in order.\n\n### heapq\n\nPython's **`heapq`** module does the same in C, on a plain list:\n\n```pycon\n>>> import heapq\n>>> heap = [5, 3, 8, 1, 4]\n>>> heapq.heapify(heap)\n>>> heap\n[1, 3, 8, 5, 4]\n>>> heapq.heappush(heap, 2)\n>>> heap[0]\n1\n>>> heapq.heappop(heap)\n1\n>>> heapq.heappop(heap)\n2\n```\n\n`heapify` turns a list into a heap in place in O(n), cheaper than n pushes, which cost O(n log n). `heappush` and `heappop` are O(log n), and `heap[0]` is the smallest, O(1). Mind two things. Only `heap[0]` has a known place: `heap[-1]` isn't the largest, and `heap[1]` isn't always the second smallest. And `heappop` trusts that the list is a heap: on a list you never heapified, it takes whatever is at index 0.\n\n`heappush`, `heappop` and `heapify` keep the smallest in front. For the largest first, push `-x` and negate what you pop, which works in every version of Python. Python 3.14 also has max-heap versions, such as `heapq.heappush_max` and `heapq.heappop_max`. Tuples compare item by item, so `(priority, name)` pairs come out by priority, with the name deciding ties.\n\n### The k largest, as they come\n\nTo keep the k largest values of a stream, keep a heap of at most k values. Its root is the smallest of the k largest so far: the **k-th largest**. When a value comes, push it, and if the heap now has k + 1 values, pop the smallest, which can't be among the k largest any more:\n\n```python\nimport heapq\n\n\ndef kth_largest(values, k):\n    heap = []\n    for value in values:\n        heapq.heappush(heap, value)\n        if len(heap) > k:\n            heapq.heappop(heap)\n    return heap[0]\n\n\nscores = [40, 95, 12, 78, 66, 88, 51]\nprint(kth_largest(scores, 3), sorted(scores)[-3])\n```\n\n```output\n78 78\n```\n\nEach value costs O(log k), so a stream of n values costs O(n log k) time and O(k) memory, however long it is. And at every moment, `heap[0]` is the answer for the values seen so far, while a sorted list of all the values seen would cost O(n) for each new value, O(n²) in all. Once the heap is full, you can save more work: a value that isn't greater than `heap[0]` can't get in, so skip it, and for one that can, `heapq.heapreplace(heap, value)` pops the smallest and pushes the value in one sift.\n\nThe same idea merges k sorted lists: keep the next item of each list in a heap, as `(value, list number, index)`, pop the smallest and push the next item of its list. That's O(N log k) for N items in all, and `heapq.merge` does it for you.\n","task":"This `push` should build a **max-heap**, a heap with the largest value at the root, where every parent is greater than or equal to its children. Fill in the blanks: the parent's index, and the comparison that stops the sifting. The program should print:\n\n```\n[12, 9, 4, 3, 7]\n12\n```\n","kind":"fill","seed":"def push(heap, value):\n    heap.append(value)\n    i = len(heap) - 1\n    while i > 0:\n        parent = [[(i - 1) // 2]]\n        if heap[i] [[<=‖<]] heap[parent]:\n            break\n        heap[i], heap[parent] = heap[parent], heap[i]\n        i = parent\n\n\nheap = []\nfor value in [3, 9, 4, 12, 7]:\n    push(heap, value)\nprint(heap)\nprint(heap[0])\n","solution":"def push(heap, value):\n    heap.append(value)\n    i = len(heap) - 1\n    while i > 0:\n        parent = (i - 1) // 2\n        if heap[i] <= heap[parent]:\n            break\n        heap[i], heap[parent] = heap[parent], heap[i]\n        i = parent\n\n\nheap = []\nfor value in [3, 9, 4, 12, 7]:\n    push(heap, value)\nprint(heap)\nprint(heap[0])\n","hints":["The parent of index `i` is the same in a max-heap as in a min-heap: `(i - 1) // 2`.","In a max-heap, sifting up stops when the new value is no greater than its parent: `heap[i] <= heap[parent]`."],"tests":[{"name":"Output","stdin":"","expect":"[12, 9, 4, 3, 7]\n12","hidden":false}],"require":[],"forbid":[],"more":[{"task":"A printer always prints the shortest job waiting. `jobs` lists what happens, in order: a positive number is a new job of that many pages, and 0 means the printer is free and prints the shortest job waiting. Write `print_order(jobs)` that returns the page counts of the jobs in the order they are printed. A 0 never comes when no job is waiting, and jobs still waiting at the end aren't printed. The hidden tests have up to 800,000 events, so finding the shortest job with `min` each time is too slow, and so is keeping the waiting jobs in a sorted list: use `heapq`.\n\nFor example, this call:\n\n```python\nprint(print_order([5, 2, 8, 0, 1, 0, 0, 3]))\n```\n\nprints (the 2 goes first, then the 1 that came after it, then the 5):\n\n```\n[2, 1, 5]\n```\n","kind":"code","seed":"import heapq\n\n\ndef print_order(jobs):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(print_order([5, 2, 8, 0, 1, 0, 0, 3]))\n","solution":"import heapq\n\n\ndef print_order(jobs):\n    waiting = []\n    printed = []\n    for pages in jobs:\n        if pages == 0:\n            printed.append(heapq.heappop(waiting))\n        else:\n            heapq.heappush(waiting, pages)\n    return printed\n\n\nif __name__ == \"__main__\":\n    print(print_order([5, 2, 8, 0, 1, 0, 0, 3]))\n","hints":["Keep the waiting jobs in a list that is a heap: `heapq.heappush(waiting, pages)` adds a job in O(log n).","For a 0, `heapq.heappop(waiting)` removes and returns the shortest job in O(log n). Append it to the list of printed jobs.","`min(waiting)` plus `waiting.remove(...)` gives the same answers, but each is O(n): with 75,000 printouts from a line of up to 150,000 jobs, that's billions of steps. A list kept sorted with `bisect.insort` is O(n) per job too, and the biggest test makes it move hundreds of thousands of items again and again."],"tests":[{"name":"print(print_order([5, 2, 8, 0, 1, 0, 0, 3]))","stdin":"","call":"print(print_order([5, 2, 8, 0, 1, 0, 0, 3]))\n","expect":"[2, 1, 5]","hidden":false},{"name":"print_order([])","stdin":"","call":"print_order([])\n","expect":"[]","hidden":true,"returns":true},{"name":"print_order([4, 4, 0, 0])","stdin":"","call":"print_order([4, 4, 0, 0])\n","expect":"[4, 4]","hidden":true,"returns":true},{"name":"print_order([9, 0, 7, 0, 3, 6, 0])","stdin":"","call":"print_order([9, 0, 7, 0, 3, 6, 0])\n","expect":"[9, 7, 3]","hidden":true,"returns":true},{"name":"print_order([10, 20, 30])","stdin":"","call":"print_order([10, 20, 30])\n","expect":"[]","hidden":true,"returns":true},{"name":"Test 6","stdin":"","call":"result = print_order([0 if i % 4 == 3 else i * 7919 % 100_003 + 1 for i in range(300_000)])\nprint(len(result), sum(result))\n","expect":"75000 1250026454","hidden":true,"big":true},{"name":"Test 7","stdin":"","call":"result = print_order([0 if i % 8 == 7 else 2_000_000 + i if i % 2 == 0 else 2_000_000 - i for i in range(800_000)])\nprint(len(result), sum(result))\n","expect":"100000 159999900000","hidden":true,"big":true}],"require":[],"forbid":[]},{"task":"A game's leaderboard shows the bar for the top k: after each new score, the k-th highest score so far. Write `bars(scores, k)` that returns a list of those bars, one for each score from the k-th score on (before that, there aren't k scores yet). The hidden tests have up to 500,000 scores, so sorting the scores again after each one is too slow, and so is keeping all of them in a sorted list: keep a heap of the k highest instead.\n\nFor example, with k = 3, after 50, 80 and 20 the bar is 20, the lowest of the three. 90 pushes it up to 50, 70 to 70, and 60 doesn't get into the top 3. So this call:\n\n```python\nprint(bars([50, 80, 20, 90, 70, 60], 3))\n```\n\nprints:\n\n```\n[20, 50, 70, 70]\n```\n","kind":"code","seed":"import heapq\n\n\ndef bars(scores, k):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(bars([50, 80, 20, 90, 70, 60], 3))\n","solution":"import heapq\n\n\ndef bars(scores, k):\n    heap = []\n    result = []\n    for score in scores:\n        if len(heap) < k:\n            heapq.heappush(heap, score)\n        elif score > heap[0]:\n            heapq.heapreplace(heap, score)\n        if len(heap) == k:\n            result.append(heap[0])\n    return result\n\n\nif __name__ == \"__main__\":\n    print(bars([50, 80, 20, 90, 70, 60], 3))\n","hints":["Keep a heap of the k highest scores so far. Its smallest, `heap[0]`, is the k-th highest: the bar.","For each score, push it. If the heap now has more than k scores, pop the smallest, which has dropped out of the top k.","Once the heap holds exactly k scores, append `heap[0]` to the result after each score. Each score costs O(log k). (A full heap can skip a score that isn't greater than `heap[0]`, and take one that is with `heapq.heapreplace`.)"],"tests":[{"name":"print(bars([50, 80, 20, 90, 70, 60], 3))","stdin":"","call":"print(bars([50, 80, 20, 90, 70, 60], 3))\n","expect":"[20, 50, 70, 70]","hidden":false},{"name":"bars([5, 1, 9], 1)","stdin":"","call":"bars([5, 1, 9], 1)\n","expect":"[5, 5, 9]","hidden":true,"returns":true},{"name":"bars([3, 3, 3, 1], 2)","stdin":"","call":"bars([3, 3, 3, 1], 2)\n","expect":"[3, 3, 3]","hidden":true,"returns":true},{"name":"bars([4, 7], 3)","stdin":"","call":"bars([4, 7], 3)\n","expect":"[]","hidden":true,"returns":true},{"name":"bars([10, 40, 30, 20, 50], 5)","stdin":"","call":"bars([10, 40, 30, 20, 50], 5)\n","expect":"[10]","hidden":true,"returns":true},{"name":"bars([8, 2, 6, 9, 1, 7, 9], 2)","stdin":"","call":"bars([8, 2, 6, 9, 1, 7, 9], 2)\n","expect":"[2, 6, 8, 8, 8, 9]","hidden":true,"returns":true},{"name":"Test 7","stdin":"","call":"result = bars([i * 7919 % 200_003 for i in range(200_000)], 100)\nprint(len(result), result[-1], sum(result))\n","expect":"199901 199903 39827062577","hidden":true,"big":true},{"name":"Test 8","stdin":"","call":"result = bars([500_000 - i for i in range(500_000)], 50)\nprint(len(result), result[-1], sum(result))\n","expect":"499951 499951 249951002401","hidden":true,"big":true},{"name":"Test 9","stdin":"","call":"result = bars(list(range(500_000)), 50)\nprint(len(result), result[-1], sum(result))\n","expect":"499951 499950 124975251225","hidden":true,"big":true}],"require":[],"forbid":[]}]}