Heaps

import heapq
 
h = [5, 1, 8, 3]
heapq.heapify(h)  # in place, O(n)
h[0]  # 1: the smallest is always at index 0
heapq.heappush(h, 0)
heapq.heappop(h), heapq.heappop(h)  # (0, 1)
sorted(h)  # [3, 5, 8]: the heap itself isn't sorted
 
g = [2, 7]
heapq.heappushpop(g, 1)  # 1: smaller than g[0], g unchanged
heapq.heapreplace(g, 1)  # 2: pops first, then pushes 1
sorted(g)  # [1, 7]

import heapq
from dataclasses import dataclass
 
scores = [70, 95, 82]
heapq.heapify_max(scores)  # 3.14+
heapq.heappush_max(scores, 88)
heapq.heappop_max(scores), scores[0]  # (95, 88)
 
neg = [-s for s in [70, 95, 82]]  # before 3.14
heapq.heapify(neg)
-heapq.heappop(neg)  # 95
 
words = ["pear", "apple", "fig"]
heapq.heapify_max(words)  # strings can't be negated
heapq.heappop_max(words)  # 'pear'
 
 
@dataclass
class Rev:  # max-heap of strings before 3.14
    s: str
 
    def __lt__(self, other: Rev) -> bool:
        return self.s > other.s
 
 
r = [Rev(w) for w in ["pear", "apple", "fig"]]
heapq.heapify(r)
heapq.heappop(r).s  # 'pear'

import heapq
from dataclasses import dataclass, field
from itertools import count
 
h: list[tuple[int, dict[str, int]]] = []
heapq.heappush(h, (1, {"id": 1}))
try:
    heapq.heappush(h, (1, {"id": 2}))  # tie on 1
except TypeError as e:
    print(e)
# '<' not supported between instances of 'dict' and 'dict'
 
tick = count()  # a unique, increasing tiebreaker
q: list[tuple[int, int, str]] = []
for prio, job in [(2, "b"), (1, "a"), (2, "c")]:
    heapq.heappush(q, (prio, next(tick), job))
[heapq.heappop(q)[2] for _ in range(3)]  # ['a', 'b', 'c']
 
 
@dataclass(order=True)
class Task:
    prio: int
    name: str = field(compare=False)  # never compared
 
 
t = [Task(3, "ship"), Task(1, "plan"), Task(1, "spec")]
heapq.heapify(t)
heapq.heappop(t).prio  # 1: the names were never compared

import heapq
 
best = {"a": 5, "b": 3}  # the current priority of each key
h = [(5, "a"), (3, "b")]
heapq.heapify(h)
best["a"] = 1  # "decrease key": record it...
heapq.heappush(h, (1, "a"))  # ...and push a fresh entry
order: list[tuple[str, int]] = []
while h:
    prio, key = heapq.heappop(h)
    if prio != best[key]:
        continue  # stale: a newer entry already came out
    order.append((key, prio))
print(order)  # [('a', 1), ('b', 3)]

import heapq
from collections.abc import Iterable
 
 
def top_k(stream: Iterable[int], k: int) -> list[int]:
    h: list[int] = []  # min-heap of the k best so far
    for x in stream:
        if len(h) < k:
            heapq.heappush(h, x)
        else:  # O(1) when x can't beat h[0]
            heapq.heappushpop(h, x)
    return sorted(h, reverse=True)
 
 
top_k([5, 1, 9, 3, 7, 9, 2], 3)  # [9, 9, 7]
list(heapq.merge([1, 4, 9], [2, 3, 10], [5]))
# [1, 2, 3, 4, 5, 9, 10]
heapq.nlargest(2, ["aa", "b", "cccc"], key=len)
# ['cccc', 'aa']

import heapq
 
 
class Median:
    def __init__(self) -> None:
        self.lo: list[int] = []  # max-heap: the smaller half
        self.hi: list[int] = []  # min-heap: the larger half
 
    def add(self, x: int) -> None:
        heapq.heappush_max(self.lo, x)  # 3.14+
        heapq.heappush(self.hi, heapq.heappop_max(self.lo))
        if len(self.hi) > len(self.lo):
            back = heapq.heappop(self.hi)
            heapq.heappush_max(self.lo, back)
 
    def median(self) -> float:
        if len(self.lo) > len(self.hi):
            return self.lo[0]
        return (self.lo[0] + self.hi[0]) / 2
 
 
m = Median()
out: list[float] = []
for x in [5, 15, 1, 3]:
    m.add(x)
    out.append(m.median())
print(out)  # [5, 10.0, 5, 4.0]

import heapq
from collections import Counter
 
nums = [3, 2, 1, 5, 6, 4]
heapq.nlargest(2, nums)[-1]  # 5: the 2nd largest
 
pts = [(1, 3), (-2, 2), (5, 8), (0, 1)]
heapq.nsmallest(2, pts, key=lambda p: p[0] ** 2 + p[1] ** 2)
# [(0, 1), (-2, 2)]
Counter("aaabbc").most_common(2)  # [('a', 3), ('b', 2)]

import heapq
 
type Entry = tuple[int, int, str]  # (prio, count, name)
q: list[Entry] = []
heapq.heappush(q, (1, 0, "plan"))
 
bad: list[dict[str, int]] = []
heapq.heappush(bad, {"a": 1})  # ❌ type error: dicts lack <

import heapq
 
 
def min_rooms(meetings: list[tuple[int, int]]) -> int:
    ends: list[int] = []  # end times of rooms in use
    for start, end in sorted(meetings):
        if ends and ends[0] <= start:
            heapq.heapreplace(ends, end)  # reuse that room
        else:
            heapq.heappush(ends, end)
    return len(ends)
 
 
print(min_rooms([(0, 30), (5, 10), (15, 20)]))  # 2
print(min_rooms([(7, 10), (2, 4)]))  # 1

import heapq
 
 
class KthLargest:
    def __init__(self, k: int, nums: list[int]) -> None:
        self.k = k
        self.h = heapq.nlargest(k, nums)
        heapq.heapify(self.h)  # a min-heap of the k largest
 
    def add(self, x: int) -> int:
        if len(self.h) < self.k:
            heapq.heappush(self.h, x)
        else:
            heapq.heappushpop(self.h, x)
        return self.h[0]
 
 
s = KthLargest(3, [4, 5, 8, 2])
print([s.add(x) for x in [3, 5, 10, 9, 4]])
# [4, 5, 5, 8, 8]

import heapq
from collections import Counter
 
 
def reorganize(s: str) -> str:
    h = [(-n, ch) for ch, n in Counter(s).items()]
    heapq.heapify(h)  # a max-heap by count, via negation
    out: list[str] = []
    held: tuple[int, str] | None = None  # used last round
    while h:
        n, ch = heapq.heappop(h)
        out.append(ch)
        if held:
            heapq.heappush(h, held)
        held = (n + 1, ch) if n + 1 < 0 else None
    return "".join(out) if len(out) == len(s) else ""
 
 
print(reorganize("aab"), repr(reorganize("aaab")))  # aba ''

import heapq
from dataclasses import dataclass, field
from itertools import count
 
 
@dataclass(order=True)
class Entry:
    prio: int
    seq: int
    task: str = field(compare=False)
    alive: bool = field(default=True, compare=False)
 
 
class PQ:
    def __init__(self) -> None:
        self.h: list[Entry] = []
        self.live: dict[str, Entry] = {}
        self.seq = count()
 
    def push(self, task: str, prio: int) -> None:
        if task in self.live:  # update = remove + add
            self.remove(task)
        entry = Entry(prio, next(self.seq), task)
        self.live[task] = entry
        heapq.heappush(self.h, entry)
 
    def remove(self, task: str) -> None:
        self.live.pop(task).alive = False  # lazy
 
    def pop(self) -> str:
        while self.h:
            entry = heapq.heappop(self.h)
            if entry.alive:
                del self.live[entry.task]
                return entry.task
        raise KeyError("pop from an empty priority queue")
 
 
pq = PQ()
pq.push("write", 3)
pq.push("test", 2)
pq.push("ship", 5)
pq.push("ship", 1)  # raise its priority
pq.remove("test")
print(pq.pop(), pq.pop())  # ship write

import heapq
 
 
def smallest_range(lists: list[list[int]]) -> list[int]:
    h = [(xs[0], i, 0) for i, xs in enumerate(lists)]
    heapq.heapify(h)
    hi = max(xs[0] for xs in lists)  # max of current heads
    best = [h[0][0], hi]
    while True:
        lo, i, j = heapq.heappop(h)
        if hi - lo < best[1] - best[0]:
            best = [lo, hi]
        if j + 1 == len(lists[i]):
            return best  # list i is used up
        nxt = lists[i][j + 1]
        hi = max(hi, nxt)
        heapq.heappush(h, (nxt, i, j + 1))
 
 
nums = [[4, 10, 15, 24, 26], [0, 9, 12, 20], [5, 18, 22, 30]]
print(smallest_range(nums))  # [20, 24]

import heapq
 
 
def max_capital(
    k: int, w: int, profits: list[int], capital: list[int]
) -> int:
    locked = sorted(zip(capital, profits))  # by cost
    ready: list[int] = []  # max-heap of unlocked profits
    i = 0
    for _ in range(k):
        while i < len(locked) and locked[i][0] <= w:
            heapq.heappush_max(ready, locked[i][1])  # 3.14+
            i += 1
        if not ready:
            break  # nothing affordable
        w += heapq.heappop_max(ready)
    return w
 
 
print(max_capital(2, 0, [1, 2, 3], [0, 1, 1]))  # 4
print(max_capital(3, 0, [1, 2, 3], [0, 1, 2]))  # 6