Heaps and Priority Queues Questions

The idea in plain words

You can read the stream only once and may keep only k values, so you must decide on the spot, for every incoming number, whether it belongs among the k largest seen so far. To decide, you only need to compare it with the smallest member of the current top k. If it is bigger, it earns a place and that smallest member is evicted. If not, it is dropped forever.

A min-heap (a binary tree stored in an array where every parent is less than or equal to its children, so the smallest value is always at index 0) gives you exactly that "smallest of the current top k" in O(1), and lets you swap it out in O(log k). That is why a min-heap, and not a max-heap, is the right tool for the k largest: the root is the gatekeeper, the weakest member of the club.

The brute force and why it is not enough

Collect everything into a list, sort descending, take the first k. That is O(n log n) time and O(n) memory, and the memory is the real problem: the stream may be larger than RAM, or never end. Python's heapq.nlargest(k, stream) already does exactly this, so in an interview you write the heap yourself and then mention that.

The approach

For each incoming value x:

1. If the heap holds fewer than k values, push x.
2. Otherwise compare x with heap[0], the smallest of the current top k. If x > heap[0], call heapq.heapreplace(heap, x), which removes the root and inserts x with a single sift (one pass down the tree) instead of a pop followed by a push.
3. Otherwise discard x.

After the stream ends, sort the (at most k) values descending. The class TopK in the code is the incremental version: insert is the same three steps, and top() returns the current top k sorted descending, so memory stays proportional to k whatever the number of inserts.

The code

import heapq
import random


def top_k_stream(stream, k, trace=False):
    """Return the k largest values of a one-pass iterable, largest first."""
    if k <= 0:
        return []
    heap = []  # min-heap: heap[0] is the smallest of the current top k
    for x in stream:
        if len(heap) < k:
            heapq.heappush(heap, x)
            action = "push (heap not full)"
        elif x > heap[0]:
            heapq.heapreplace(heap, x)  # pop the root and push x in one sift
            action = "replace root"
        else:
            action = "discard"
        if trace:
            print(f"read {x:>3}: {action:<21} heap={sorted(heap)}")
    return sorted(heap, reverse=True)


class TopK:
    """Incremental version: insert values, read the current top k."""

    def __init__(self, k):
        self.k = k
        self._heap = []

    def insert(self, x):
        if len(self._heap) < self.k:
            heapq.heappush(self._heap, x)
        elif self._heap and x > self._heap[0]:
            heapq.heapreplace(self._heap, x)

    def top(self):
        return sorted(self._heap, reverse=True)  # O(k log k), memory stays O(k)


def numbers():
    yield from [5, 1, 9, 3, 7, 7, 2, 8, 6, 4]


print(top_k_stream(numbers(), 3, trace=True))

# k larger than the stream, k = 0, ties
print(top_k_stream(iter([4, 2]), 5), top_k_stream(iter([4, 2]), 0), top_k_stream(iter([7, 7, 7, 1]), 2))

# incremental API, for example the slowest requests seen so far
slowest = TopK(3)
for ms in [120, 80, 450, 95, 300, 510, 40]:
    slowest.insert(ms)
    print("after", ms, "->", slowest.top())

# check against sorting everything, on random streams with duplicates
rng = random.Random(7)
for _ in range(500):
    data = [rng.randint(-20, 20) for _ in range(rng.randint(0, 30))]
    k = rng.randint(0, 8)
    assert top_k_stream(iter(data), k) == sorted(data, reverse=True)[:k]
print("500 random streams match sorted(...)[:k]")

Output:

read   5: push (heap not full)  heap=[5]
read   1: push (heap not full)  heap=[1, 5]
read   9: push (heap not full)  heap=[1, 5, 9]
read   3: replace root          heap=[3, 5, 9]
read   7: replace root          heap=[5, 7, 9]
read   7: replace root          heap=[7, 7, 9]
read   2: discard               heap=[7, 7, 9]
read   8: replace root          heap=[7, 8, 9]
read   6: discard               heap=[7, 8, 9]
read   4: discard               heap=[7, 8, 9]
[9, 8, 7]
[4, 2] [] [7, 7]
after 120 -> [120]
after 80 -> [120, 80]
after 450 -> [450, 120, 80]
after 95 -> [450, 120, 95]
after 300 -> [450, 300, 120]
after 510 -> [510, 450, 300]
after 40 -> [510, 450, 300]
500 random streams match sorted(...)[:k]

Worked trace

The trace lines above come from top_k_stream(numbers(), 3, trace=True), with k = 3. The heap contents are printed sorted for readability; the root is the first (smallest) value shown.

• Reads 5, 1, 9 fill the heap, so it holds [1, 5, 9] with root 1.
• Read 3 beats the root 1, so 1 is evicted: [3, 5, 9].
• Read 7 beats root 3: [5, 7, 9]. The second 7 beats root 5: [7, 7, 9].
• Read 2 is below the root 7, so it is discarded and the heap is unchanged.
• Read 8 beats root 7: [7, 8, 9]. Reads 6 and 4 are discarded.
• The final sort descending gives [9, 8, 7].

The TopK run shows the same logic as a live read-out: after 95 arrives the three slowest requests are [450, 120, 95], and the later 40 changes nothing.

Complexity

• Time: each of the n values costs one comparison with the root, and at most one O(log k) heappush or heapreplace. Total O(n log k), plus O(k log k) for the final sort. Since k is at most n, this is never worse than sorting everything, and when k is small it is close to O(n).
• Space: the heap never holds more than k values, so O(k) extra memory. The returned list is also size k. Nothing proportional to n is stored.
• TopK.top() costs O(k log k) per call because it sorts a copy. If reads are much more frequent than inserts, cache the sorted list and rebuild it only after an insert that actually changed the heap.

Edge cases and common mistakes

• k larger than the stream: the heap never fills, so you get every value, sorted (the [4, 2] line).
• k = 0 (or negative): return an empty list immediately. Without that guard, heap[0] on an empty heap raises IndexError.
• Ties: with x > heap[0] an equal value is discarded. For plain numbers the returned values are the same either way, because equal numbers are interchangeable. If equal keys belong to different items (requests with the same latency, say), which of them stays is not determined by the comparison alone: switching to >= does not give a "newest wins" rule either, since it replaces whichever equal item happens to sit at the root, not the oldest one. To get a definite rule, push (key, arrival_number) pairs: pairs never tie, so the k largest pairs are exactly the largest keys with the newest items winning among equal keys.
• Empty stream: returns [].
• A frequent mistake is a max-heap of size n or k. A max-heap of size k puts the largest value at the root, which tells you nothing about who to evict.
• Another is heappop then heappush on every replacement. It works, but does two sifts where heapreplace does one.
• Up to Python 3.13, heapq provides min-heap functions only (Python 3.14 added max-heap counterparts such as heappush_max, heappop_max and heapify_max), and a min-heap is exactly what is needed here. For the k smallest you would negate the values, flip the comparison, or on 3.14 use the max-heap functions.

Running the code

python3 top_k.py

No packages are needed beyond the standard library (heapq, random); the file runs unchanged on Python 3.12.

The idea in plain words

The base structure is the two-heap running median: the lower half of the items seen so far sits in a max-heap (largest on top), the upper half in a min-heap (smallest on top), and their sizes are equal or the lower half has one more. Each insert pushes the item into the half it belongs to, then moves one top across if the sizes drifted; the median is read from the tops (the lower top for an odd count, the average of both tops for an even count).

That structure never needs arithmetic on the values, only an ordering: which of two items comes first. So to support objects, strings or timestamps, give the heaps an ordering key and make three decisions: how a max-heap is built when negation does not exist, how ties are broken, and what the median means when two middle items cannot be averaged. The key and the tie-break are needed for every type of value, the even-count rule depends on the type, and invalid values (None, NaN) are stopped at the door before they reach a heap.

How to adapt the heaps

• Key function. The structure takes key(obj), for example lambda r: r["score"], and orders objects by that key. The default is the object itself (numbers, strings, datetimes).
• A max-heap without negation. Up to Python 3.13 the heapq functions all work on a min-heap, and the usual trick of storing -x only works for numbers. For strings or datetimes wrap the key in a small class whose < is reversed (Desc in the code), so the min-heap pops the largest key first; this wrapper works on every Python version, and the code below uses it. Python 3.14 added max-heap counterparts (heapify_max, heappush_max, heappop_max, heapreplace_max, heappushpop_max) that keep the largest item at index 0 without any reversing wrapper; code that must run on older versions still needs the wrapper. For plain numbers -x is simpler and faster than the wrapper, because the wrapper adds a Python-level __lt__ call to every comparison.
• Entries are (key, seq, obj). The second field, seq, is a unique insertion counter, explained next.

Ties and duplicates

When two entries have equal keys, a tuple comparison goes on to the next field. Without a counter that next field is the object itself, and for dicts or custom objects the comparison raises TypeError (the run shows '<' not supported between instances of 'dict' and 'dict'). A unique seq means two entries never tie, the objects are never compared, and the order is a total order (a rule that says, for any two entries, which comes first, with no two entries tied): (key, insertion order). Equal keys are kept in arrival order, and "the median object" is then well defined: the middle element of the list sorted on (key, insertion order). Two details make this hold:

• The new item goes to the lower half only if its key is strictly smaller than the lower half's top; an equal key goes to the upper half, because the new item is the newest and so is larger in the total order.
• The lower half is a max-heap on the total order, so among equal keys its top must be the newest. It therefore stores -seq, and Desc defines __eq__ as well as __lt__ (the special methods Python calls for == and <): tuple comparison tests equality first, and without __eq__ it would never get to the seq field.

A tie walked through, with the scores 70, 90, 70 for ann, bob, cy. After ann the lower half holds (70, -1, ann). bob has key 90, not below 70, so it goes to the upper half. cy has key 70, not strictly below the lower top's 70, so it goes to the upper half too, which then has two entries and one more than the lower half; its smallest entry, (70, 3, cy), moves down and is stored as (70, -3, cy). The lower half now holds (70, -3, cy) and (70, -1, ann), and the smaller second field, -3, belongs to the newer entry, so cy is on top. That is right: sorted on (key, insertion order) the three are ann, cy, bob, and the middle is cy, matching the line added cy score 70: middle objects cy, cy in the output. With +seq the older ann would be on top and the answer would be wrong. Without __eq__ on Desc, two different Desc objects with the same key are never "equal" to the tuple comparison, so it never looks at seq, and which of equal entries ends up on top depends on the order they were pushed in (checked with three equal-key entries pushed in opposite orders, which put different entries on top).

Even counts: what is the median of two objects?

For an even number of items there are two middle objects. middle_pair() returns both (the same object twice for odd counts). The median of the keys depends on the type:

• Numbers: the mean of the two keys. Scores 70, 90, 70, 85, 70, 95 have keys sorted 70, 70, 70, 85, 90, 95, so the median is (70 + 85) / 2 = 77.5, as the last line of the first block of output shows.
• Datetimes: the midpoint lo + (hi - lo) / 2, because subtracting two datetimes gives a timedelta (a length of time, such as 2 hours 30 minutes) and halving it is defined.
• Strings and arbitrary objects: there is no midpoint. Return the lower middle (or the pair) and say which. The code returns the lower middle key.

Strings

Python compares strings by Unicode code point (each character's numeric id, which gives dictionary-like ordering letter by letter, called lexicographic order), so uppercase letters sort before lowercase: 'Zed' < 'apple'. The output shows ['Apple', 'Zed', 'apple', 'fig', 'pear'] and the median string 'apple'. With key=str.casefold (a lowercasing that is meant for case-insensitive comparison) the order becomes Apple, apple, fig, pear, Zed and the median is 'fig'. Choose the key to match the requirement; for language-aware ordering (accents and alphabets sorted as a reader of that language expects) use a locale-aware key such as locale.strxfrm, which converts a string into a form whose plain comparison follows the current locale's rules.

Timestamps

Parse at the boundary into timezone-aware datetime objects (ones that carry their offset from UTC, as opposed to naive ones that carry none) (here datetime.fromisoformat as the key) so the ordering is on the instant, not on the text: 12:30+02:00 is 10:30 UTC, which sorts after 10:00+00:00. Mixing naive and aware datetimes raises TypeError: can't compare offset-naive and offset-aware datetimes; the structure compares the new key with a stored one before it changes anything, so the failed insert leaves the heaps untouched. An alternative is to convert to integer epoch seconds (the count of seconds since 1 January 1970 UTC) on entry and keep numbers inside.

NaN and None

None has no order against numbers and raises TypeError when compared. NaN (not a number) is worse: every comparison with it is false, so it does not raise, and the heap silently loses its order. The run pushes 3.0, nan, 1.0, 2.0 into a plain heap and pops [1.0, 2.0, nan, 3.0], which is not sorted, so any median read from such a heap is wrong. The structure therefore rejects None and NaN up front (k is None or k != k, since NaN is the only value not equal to itself) with ValueError. A string offered to a structure of numbers raises TypeError during the first comparison, before any mutation; the last line of that block shows the structure intact. Decide the policy once: reject, or skip and count. Do not let an invalid value into the heaps.

The code

import heapq
import numbers
import random
from datetime import datetime


class Desc:
    """Wraps a key so that a min-heap pops the LARGEST key first (works for strings, datetimes, anything orderable)."""

    __slots__ = ("key",)

    def __init__(self, key):
        self.key = key

    def __lt__(self, other):
        return other.key < self.key

    def __eq__(self, other):                         # tuple comparison tests == first; without this it never reaches seq
        return self.key == other.key


class RunningMedian:
    """Running median of arbitrary objects ordered by key(obj).

    Heap entries are (key, seq, obj). seq is a unique insertion counter, so two entries never tie and the
    objects themselves are never compared. The order is (key, insertion order), a total order. The lower half is a
    max-heap on that order, which is why it stores -seq: among equal keys the newest entry must come out first."""

    def __init__(self, key=lambda obj: obj):
        self.key = key
        self.small = []            # lower half: (Desc(key), -seq, obj); -seq pops the newest of equal keys first
        self.large = []            # upper half: (key, seq, obj)
        self.seq = 0

    def add(self, obj):
        k = self.key(obj)
        if k is None or k != k:                          # None, or NaN (the only value that is not equal to itself)
            raise ValueError(f"cannot order {k!r}")
        # Compare BEFORE changing anything: an incomparable key raises TypeError and leaves the heaps intact.
        goes_low = not self.small or k < self.small[0][0].key
        self.seq += 1
        if goes_low:
            heapq.heappush(self.small, (Desc(k), -self.seq, obj))
        else:
            heapq.heappush(self.large, (k, self.seq, obj))
        if len(self.small) > len(self.large) + 1:
            d, neg_seq, o = heapq.heappop(self.small)
            heapq.heappush(self.large, (d.key, -neg_seq, o))
        elif len(self.large) > len(self.small):
            k2, s, o = heapq.heappop(self.large)
            heapq.heappush(self.small, (Desc(k2), -s, o))

    def middle_pair(self):
        """(lower middle object, upper middle object); the same object twice for an odd count."""
        if not self.small:
            raise ValueError("median of an empty stream")
        lower = self.small[0][2]
        return lower, (lower if len(self.small) > len(self.large) else self.large[0][2])

    def median_key(self):
        """Median of the keys where an average exists (numbers, datetimes); otherwise the lower middle key."""
        lo, hi = (self.key(o) for o in self.middle_pair())
        if lo == hi:
            return lo
        if isinstance(lo, (numbers.Real, datetime)):
            return lo + (hi - lo) / 2
        return lo


def show(title, fn):
    print(f"-- {title}")
    try:
        fn()
    except Exception as e:                                # the demo prints the failure it provokes
        print(f"   {type(e).__name__}: {e}")


if __name__ == "__main__":
    # 1. Objects ordered by a numeric field, with tied scores
    m = RunningMedian(key=lambda r: r["score"])
    for name, score in [("ann", 70), ("bob", 90), ("cy", 70), ("di", 85), ("ed", 70), ("flo", 95)]:
        m.add({"name": name, "score": score})
        lo, hi = m.middle_pair()
        print(f"added {name:<3} score {score}: middle objects {lo['name']}, {hi['name']}; median score {m.median_key()}")

    # 2. Strings: Python compares code points, so upper case sorts before lower case
    words = ["pear", "Apple", "fig", "apple", "Zed"]
    print("\nsorted by code point:", sorted(words), " casefolded:", sorted(words, key=str.casefold))
    for label, key in (("code point", lambda w: w), ("casefold", str.casefold)):
        m = RunningMedian(key=key)
        for w in words:
            m.add(w)
        print(f"median string by {label}: {m.middle_pair()[0]!r}")

    # 3. Timestamps: parse to timezone-aware datetimes at the boundary
    m = RunningMedian(key=datetime.fromisoformat)
    for t in ["2026-03-01T10:00:00+00:00", "2026-03-01T12:30:00+02:00", "2026-03-01T09:00:00+00:00", "2026-03-02T00:00:00+00:00"]:
        m.add(t)
    print("\nmiddle timestamps:", m.middle_pair(), "-> midpoint", m.median_key().isoformat())
    show("naive and aware timestamps mixed", lambda: m.add("2026-03-01T10:00:00"))

    # 4. Why the tiebreak counter exists
    def tie_without_counter():
        h = []
        heapq.heappush(h, (70, {"name": "ann"}))
        heapq.heappush(h, (70, {"name": "cy"}))
    show("equal keys, objects compared (no counter)", tie_without_counter)

    # 5. NaN and None
    h = []
    for v in [3.0, float("nan"), 1.0, 2.0]:
        heapq.heappush(h, v)
    print("\nplain heap with NaN pops in this order:", [heapq.heappop(h) for _ in range(len(h))], "(not sorted)")
    m = RunningMedian()
    show("NaN rejected", lambda: m.add(float("nan")))
    show("None rejected", lambda: m.add(None))
    m.add(2); m.add(1)
    show("string into a numeric structure", lambda: m.add("x"))
    print("structure still intact after the rejections:", m.middle_pair(), m.median_key())

    # 6. Check against a sort on (key, insertion order)
    rng = random.Random(8)
    for _ in range(3000):
        m, items = RunningMedian(key=lambda r: r[0]), []
        for i in range(rng.randint(1, 25)):
            item = (rng.randint(0, 4), i)                 # (score, id): many equal scores
            m.add(item); items.append(item)
            ordered = sorted(items, key=lambda r: (r[0], r[1]))
            n = len(ordered)
            want = (ordered[(n - 1) // 2], ordered[n // 2])
            assert m.middle_pair() == want, (items, m.middle_pair(), want)
    print("\n3000 random streams with heavy ties: middle objects equal the sort on (key, insertion order)")

Run, it prints:

added ann score 70: middle objects ann, ann; median score 70
added bob score 90: middle objects ann, bob; median score 80.0
added cy  score 70: middle objects cy, cy; median score 70
added di  score 85: middle objects cy, di; median score 77.5
added ed  score 70: middle objects ed, ed; median score 70
added flo score 95: middle objects ed, di; median score 77.5

sorted by code point: ['Apple', 'Zed', 'apple', 'fig', 'pear']  casefolded: ['Apple', 'apple', 'fig', 'pear', 'Zed']
median string by code point: 'apple'
median string by casefold: 'fig'

middle timestamps: ('2026-03-01T10:00:00+00:00', '2026-03-01T12:30:00+02:00') -> midpoint 2026-03-01T10:15:00+00:00
-- naive and aware timestamps mixed
   TypeError: can't compare offset-naive and offset-aware datetimes
-- equal keys, objects compared (no counter)
   TypeError: '<' not supported between instances of 'dict' and 'dict'

plain heap with NaN pops in this order: [1.0, 2.0, nan, 3.0] (not sorted)
-- NaN rejected
   ValueError: cannot order nan
-- None rejected
   ValueError: cannot order None
-- string into a numeric structure
   TypeError: '<' not supported between instances of 'str' and 'int'
structure still intact after the rejections: (1, 2) 1.5

3000 random streams with heavy ties: middle objects equal the sort on (key, insertion order)

Complexity and pitfalls

• Insert is O(log n) comparisons and the median read is O(1), as for plain numbers. The constant is higher: each lower-half entry allocates a Desc wrapper, and every comparison calls Python code.
• Space is O(n) entries of three fields.
• Pitfall: comparing the objects when keys tie. Use a unique counter.
• Pitfall: using -key on strings or datetimes: TypeError.
• Pitfall: mixing numeric and string keys, naive and aware datetimes.
• Pitfall: letting NaN in. It corrupts the order without an error.
• The final check in the run compares the middle objects with a sort on (key, insertion order) over 3000 random streams with five distinct scores, so equal keys occur constantly.

Running the code

python generic_median.py

Python 3.12, standard library only (heapq, numbers, random, datetime). No install or config needed.

The idea

A two-heap median works only if two rules (invariants: facts that must be true after every operation) hold: sizes stay within one of each other, and every value in lowers is less than or equal to every value in highers. Checking sizes alone cannot protect the second rule. If the code decides which heap gets a new value by looking only at the sizes, then a new value can land in the wrong half. That is the first bug to suspect when "the top of lowers is greater than the top of highers".

The failing input and the wrong behaviour

The review code decides by size only: when lowers is not larger it pushes there, otherwise it pushes to highers. Feed it 9 then 8: 9 goes to lowers (sizes tied), then 8 goes to highers (because lowers is larger), so lowers has top 9 and highers has top 8. The ordering rule is broken after two inserts. Feed 1, 2, 3 and the reported median is 3, but the true median is 2 (output below). The 9, 8 case still reports the right median, 8.5, by luck, which is why this bug can survive casual tests: the violation is latent until a later insert makes the wrong top the median.

Debugging steps

1. Reproduce small. Generate random short streams (here values 1 to 9, up to 8 long), compare every step against statistics.median of the prefix, and keep the shortest failing stream. A two-element repro is far easier to reason about than a log of thousands.
2. Assert the invariants after every insert, not only at the end, so the failure points at the first operation that breaks a rule. check_invariants below checks both rules and names which one broke.
3. Print the tops and sizes of both heaps after each add on the small repro and read the state transition that went wrong.
4. Ask what decides placement. If the answer is "only the sizes", that is the cause. If the code does compare values, check which heap each comparison reads and the sign (see the second cause below).

Fix

Route through the lower half and move its largest value across. In steps: (1) push the new value into lowers; (2) pop the largest value now in lowers and push it into highers; (3) if highers now has more items than lowers, pop its smallest and push it back into lowers. An alternative fix is to compare the new value with the top of lowers: push to lowers if it is less than or equal, else highers, then rebalance sizes. Both are O(log n) per insert.

Why ordering holds. The value sent to highers in step 2 is the largest of everything in lowers plus the new value, so whatever stays in lowers is at most that value, and the old contents of highers were already at least everything in lowers. Step 3 moves the smallest of highers back, which is at least everything left in lowers. On numbers, here is the same three-step add with a return value that reports what moved, run on the stream 9, 8, 1:

import heapq


class FixedMedian:
    def __init__(self):
        self.lowers = []    # max-heap, stored negated
        self.highers = []   # min-heap

    def add(self, x):
        heapq.heappush(self.lowers, -x)
        moved = -heapq.heappop(self.lowers)
        heapq.heappush(self.highers, moved)
        back = None
        if len(self.highers) > len(self.lowers):
            back = heapq.heappop(self.highers)
            heapq.heappush(self.lowers, -back)
        return moved, back


m = FixedMedian()
for x in [9, 8, 1]:
    moved, back = m.add(x)
    print(f"add {x}: moved {moved} to highers, moved back {back}; "
          f"lowers={sorted(-v for v in m.lowers)} highers={sorted(m.highers)}")

Output:

add 9: moved 9 to highers, moved back 9; lowers=[9] highers=[]
add 8: moved 9 to highers, moved back None; lowers=[8] highers=[9]
add 1: moved 8 to highers, moved back 8; lowers=[1, 8] highers=[9]

After add 8 the tops are 8 and 9, in order, whereas the size-only code left 9 above 8. After add 1 the largest of {8, 1} goes up, highers is briefly bigger ([8, 9] against [1]), so the smallest of it, 8, comes back; the sizes are 2 and 1 and the tops are 8 and 9.

import heapq
import random
import statistics


class BuggyMedian:
    """Balances by size only: it never compares the new value with either top."""

    def __init__(self):
        self.lowers = []    # max-heap, stored negated
        self.highers = []   # min-heap

    def add(self, x):
        if len(self.lowers) <= len(self.highers):
            heapq.heappush(self.lowers, -x)
        else:
            heapq.heappush(self.highers, x)

    def median(self):
        if len(self.lowers) > len(self.highers):
            return -self.lowers[0]
        return (-self.lowers[0] + self.highers[0]) / 2


class FixedMedian(BuggyMedian):
    def add(self, x):
        # route by value through lowers, then move its largest across: ordering holds by construction
        heapq.heappush(self.lowers, -x)
        heapq.heappush(self.highers, -heapq.heappop(self.lowers))
        if len(self.highers) > len(self.lowers):             # keep lowers equal or one larger
            heapq.heappush(self.lowers, -heapq.heappop(self.highers))


def check_invariants(m):
    assert len(m.lowers) in (len(m.highers), len(m.highers) + 1), "size invariant broken"
    if m.lowers and m.highers:
        assert -m.lowers[0] <= m.highers[0], f"ordering broken: top(lowers)={-m.lowers[0]} > top(highers)={m.highers[0]}"


def find_failure(cls, trials=2000):
    """Smallest random input (by length) on which cls breaks an invariant or returns a wrong median."""
    random.seed(11)
    worst = None
    for _ in range(trials):
        xs = [random.randint(1, 9) for _ in range(random.randint(1, 8))]
        m = cls()
        for i, x in enumerate(xs):
            m.add(x)
            try:
                check_invariants(m)
                assert m.median() == statistics.median(xs[: i + 1])
            except AssertionError as e:
                if worst is None or i + 1 < len(worst[0]):
                    worst = (xs[: i + 1], str(e) or "wrong median")
                break
    return worst


stream = [1, 2, 3]
for cls in (BuggyMedian, FixedMedian):
    m = cls()
    for x in stream:
        m.add(x)
    print(f"{cls.__name__}: stream={stream} lowers={sorted(-v for v in m.lowers)} highers={sorted(m.highers)} median={m.median()}")

print("buggy, smallest failing input found:", find_failure(BuggyMedian))
print("fixed, failing input found:", find_failure(FixedMedian))

Output:

BuggyMedian: stream=[1, 2, 3] lowers=[1, 3] highers=[2] median=3
FixedMedian: stream=[1, 2, 3] lowers=[1, 2] highers=[3] median=2
buggy, smallest failing input found: ([9, 8], 'ordering broken: top(lowers)=9 > top(highers)=8')
fixed, failing input found: None

Reading the harness. check_invariants asserts the two rules on one object: the size rule (lowers has as many items as highers, or one more) and the ordering rule (the largest of lowers, which is -m.lowers[0] because values are stored negated, must not exceed highers[0]). find_failure seeds the random generator so every run is the same, then tries 2000 random streams of 1 to 8 values from 1 to 9. For each stream it adds the values one at a time and, after each add, runs check_invariants and compares m.median() with statistics.median of the values seen so far. Any AssertionError is caught: the code keeps the shortest failing prefix in worst as a pair of (the prefix, the message) and abandons that stream with break. It returns None when nothing ever failed. The str(e) or "wrong median" part labels a failure that came from the median comparison, whose assertion has no message.

For the stream 1, 2, 3 in the buggy class: add 1 goes to lowers (sizes 0 and 0, so lowers is not larger); add 2 goes to highers (now lowers is larger); add 3 goes to lowers again (sizes tied at 1 and 1), giving lowers=[1, 3], highers=[2]. Both sizes look legal, but 3 sits in the lower half above 2. With lowers larger, median() reads -self.lowers[0], which is 3, not the true middle value 2.

The buggy class breaks on the stream 9, 8 (two inserts); the fixed class survives 2000 random streams with both invariants checked at every step and every median compared against statistics.median.

A second cause: swapped or un-negated heaps

Up to Python 3.13, heapq offers only min-heap functions (3.14 adds heappush_max and related functions), so on 3.12 a "max-heap" must store negated values. If the code forgets the negation, the heap named lowers is really a min-heap and its top is the smallest value, not the largest, so it holds the wrong half while looking balanced. This variant keeps the same routing as the correct code and drops only the negation:

import heapq


class MissingNegation:
    """Same routing as the correct version, but `lowers` is a plain min-heap: the negation was forgotten."""

    def __init__(self):
        self.lowers, self.highers = [], []

    def add(self, x):
        heapq.heappush(self.lowers, x)                          # should be -x
        heapq.heappush(self.highers, heapq.heappop(self.lowers))  # should negate the popped value
        if len(self.highers) > len(self.lowers):
            heapq.heappush(self.lowers, heapq.heappop(self.highers))

    def median(self):
        if len(self.lowers) > len(self.highers):
            return self.lowers[0]
        return (self.lowers[0] + self.highers[0]) / 2


m = MissingNegation()
for x in [1, 2, 3, 4]:
    m.add(x)
    print(f"add {x}: lowers={sorted(m.lowers)} highers={sorted(m.highers)} top(lowers)={m.lowers[0]} top(highers)={m.highers[0] if m.highers else None}")
print("median reported:", m.median(), "(true median of 1, 2, 3, 4 is 2.5)")

add 1: lowers=[1] highers=[] top(lowers)=1 top(highers)=None
add 2: lowers=[2] highers=[1] top(lowers)=2 top(highers)=1
add 3: lowers=[1, 3] highers=[2] top(lowers)=1 top(highers)=2
add 4: lowers=[3, 4] highers=[1, 2] top(lowers)=3 top(highers)=1
median reported: 2.0 (true median of 1, 2, 3, 4 is 2.5)

Stepping through it by hand. The class pushes into lowers and then pops from it, but heappop always returns the smallest value, so what moves to highers is the smallest item, not the largest. Add 1: push 1 into lowers, pop it into highers, and since highers is now bigger, pop it back: lowers=[1], highers=[]. Add 2: lowers becomes {1, 2}, the pop sends 1 to highers, sizes are 1 and 1: lowers=[2], highers=[1], and the tops are already crossed (2 above 1). Add 3: lowers becomes {2, 3}, the pop sends 2 across, highers is {1, 2} and is bigger, so its smallest, 1, comes back: lowers=[1, 3], highers=[2]. Add 4: lowers becomes {1, 3, 4}, the pop sends 1 across, highers is {1, 2} and the sizes are 2 and 2: lowers=[3, 4], highers=[1, 2]. That is the 3 and the 1 in the last printed line.

After 4 inserts the heap named lowers has top 3 while highers has top 1, and the median it reports is 2.0 instead of 2.5. An ordering check catches it on the first call where the tops cross, but it must read the stored values the way this class stores them: m.lowers[0] <= m.highers[0], with no minus sign. check_invariants above negates lowers[0] because the correct class stores negated values, so it would not fire here. The plain form fails after add 2, when lowers[0] is 2 and highers[0] is 1. A related slip is swapping the two heaps in the median read, which makes the answer the wrong middle value. In both variants the signal is the same: the tops are out of order or the median is off by one rank.

Complexity and prevention

Insert is O(log n) and the invariant check is O(1) per call, so it is cheap to leave on in tests. Prevent regressions with the randomized comparison against a sorted reference, plus a test where the stream is decreasing (9, 8, 7, ...), which trips size-only routing at once, and one with duplicates, which pins down the tie rule: a new value equal to the top of lowers goes to lowers, and the ordering check uses less-than-or-equal.

Running the code

python3 median_bug_hunt.py
python3 missing_negation.py

Python 3.12, standard library only.

The idea in plain words

Dijkstra's algorithm repeatedly takes the unsettled node with the smallest known distance, then tries to improve its neighbours. Whenever a neighbour gets a shorter distance, its position in the queue must change: that is decrease_key. A standard binary heap (the array-based tree where each parent is no larger than its children) cannot do this: it does not know where an item sits. Notation used below: V is the number of nodes (vertices) and E the number of edges. An edge weight is the cost of travelling along that edge. To relax an edge u to v of weight w means to check whether going through u gives v a shorter route (dist[u] + w < dist[v]) and, if so, to record the better distance. A node is settled once it has been popped with its final distance. A graph is sparse when E is only a small multiple of V and dense when E approaches V squared. The graph is stored as an adjacency list: for each node, a list of (neighbour, weight) pairs.

There are two ways to get the effect.

1. Indexed heap. Keep a dictionary pos from each item to its index in the array, updated on every swap. decrease_key(item, p) looks up the index in O(1), lowers the priority, and sifts the entry up (swaps it with its parent while it is smaller) in O(log n).
2. Duplicates with skipping (lazy). Never change an entry. On an improvement, push a new entry (new_distance, node). When an entry is popped whose distance is larger than the node's current best distance, it is stale: skip it.

The queue design

The indexed queue has push(priority, item), pop() and decrease_key(item, new_priority). The heap is an array of [priority, item] pairs (lists rather than tuples, so decrease_key can overwrite the priority in place); every swap writes both new positions into pos, and pop deletes the popped item from pos. An item is "in the queue" exactly when it is a key of pos. The caller does one of three things on an improvement: insert (never seen), decrease_key (in the queue), or ignore (already settled, which cannot improve when weights are non-negative).

The lazy queue is the same array heap without pos; the algorithm handles staleness itself.

The code

The program counts heap comparisons for both variants (the two share one heap class), traces both on a 4-node graph, verifies both against Bellman-Ford (a slower but very simple algorithm that relaxes every edge V times; it serves only as a trusted reference), and measures three graphs: sparse random, dense random, and a staircase built so that node j is improved once by every earlier node.

In the code, dist.get(v, float("inf")) reads the best known distance to v, or infinity when v has not been reached yet, so the first discovery always counts as an improvement.

import random


class MinHeap:
    """Binary min-heap of [priority, item]. With track_positions=True it also records where each item sits."""

    def __init__(self, track_positions):
        self.a = []
        self.pos = {} if track_positions else None   # item -> index in self.a
        self.comparisons = 0
        self.peak = 0

    def __len__(self):
        return len(self.a)

    def _less(self, i, j):
        self.comparisons += 1
        return self.a[i] < self.a[j]

    def _swap(self, i, j):
        a = self.a
        a[i], a[j] = a[j], a[i]
        if self.pos is not None:
            self.pos[a[i][1]] = i
            self.pos[a[j][1]] = j

    def _sift_up(self, i):
        while i > 0 and self._less(i, (i - 1) // 2):
            self._swap(i, (i - 1) // 2)
            i = (i - 1) // 2

    def _sift_down(self, i):
        n = len(self.a)
        while True:
            m = i
            for c in (2 * i + 1, 2 * i + 2):
                if c < n and self._less(c, m):
                    m = c
            if m == i:
                return
            self._swap(i, m)
            i = m

    def push(self, priority, item):
        self.a.append([priority, item])
        if self.pos is not None:
            self.pos[item] = len(self.a) - 1
        self._sift_up(len(self.a) - 1)
        self.peak = max(self.peak, len(self.a))

    def pop(self):
        top = self.a[0]
        last = self.a.pop()
        if self.pos is not None:
            del self.pos[top[1]]
        if self.a:
            self.a[0] = last
            if self.pos is not None:
                self.pos[last[1]] = 0
            self._sift_down(0)
        return top

    def decrease_key(self, item, new_priority):
        i = self.pos[item]
        assert new_priority < self.a[i][0]
        self.a[i][0] = new_priority
        self._sift_up(i)


def dijkstra_lazy(graph, source, trace=False):
    """Duplicates: push a new entry on every improvement, skip entries whose distance is out of date."""
    dist = {source: 0}
    heap = MinHeap(track_positions=False)
    heap.push(0, source)
    pushes, stale = 1, 0
    while len(heap):
        d, u = heap.pop()
        if trace:
            print(f"  pop ({d}, node {u})" + ("  -> stale, skipped" if d > dist[u] else ""))
        if d > dist[u]:
            stale += 1                                # an older, worse entry for u
            continue
        for v, w in graph[u]:
            nd = d + w
            if nd < dist.get(v, float("inf")):
                dist[v] = nd
                heap.push(nd, v)
                pushes += 1
                if trace:
                    print(f"  push ({nd}, node {v}); heap now {sorted(map(tuple, heap.a))}")
    return dist, {"pushes": pushes, "stale pops": stale, "peak size": heap.peak, "comparisons": heap.comparisons}


def dijkstra_indexed(graph, source, trace=False):
    """One entry per node: insert on first discovery, decrease_key on every later improvement."""
    dist = {source: 0}
    done = set()
    heap = MinHeap(track_positions=True)
    heap.push(0, source)
    pushes, decreases = 1, 0
    while len(heap):
        d, u = heap.pop()
        done.add(u)
        if trace:
            print(f"  pop ({d}, node {u})")
        for v, w in graph[u]:
            nd = d + w
            if v not in done and nd < dist.get(v, float("inf")):
                if v in heap.pos:
                    heap.decrease_key(v, nd)
                    decreases += 1
                    if trace:
                        print(f"  decrease node {v} to {nd}; heap now {sorted(map(tuple, heap.a))}")
                else:
                    heap.push(nd, v)
                    pushes += 1
                    if trace:
                        print(f"  push ({nd}, node {v}); heap now {sorted(map(tuple, heap.a))}")
                dist[v] = nd
    return dist, {"pushes": pushes, "decrease-keys": decreases, "peak size": heap.peak, "comparisons": heap.comparisons}


def bellman_ford(graph, source):
    dist = {source: 0}
    for _ in range(len(graph)):
        for u in graph:
            if u in dist:
                for v, w in graph[u]:
                    if dist[u] + w < dist.get(v, float("inf")):
                        dist[v] = dist[u] + w
    return dist


def random_graph(nodes, edges, seed):
    rng = random.Random(seed)
    g = {u: [] for u in range(nodes)}
    for _ in range(edges):
        g[rng.randrange(nodes)].append((rng.randrange(nodes), rng.randint(1, 20)))
    return g


def staircase_graph(n):
    """Node j is improved once by every earlier node: the worst case for duplicates."""
    g = {u: [] for u in range(n)}
    for i in range(n - 1):
        g[i].append((i + 1, 1))
        for j in range(i + 2, n):
            g[i].append((j, 3 * (n - i)))
    return g


if __name__ == "__main__":
    small = {0: [(1, 4), (2, 1)], 1: [(3, 1)], 2: [(1, 2), (3, 5)], 3: []}
    print("lazy duplicates on a 4-node graph:")
    print(" ", dijkstra_lazy(small, 0, trace=True))
    print("indexed heap on the same graph:")
    print(" ", dijkstra_indexed(small, 0, trace=True))
    print()
    for seed in range(300):
        rng = random.Random(seed)
        g = random_graph(rng.randint(2, 12), rng.randint(1, 40), seed)
        want = bellman_ford(g, 0)
        assert dijkstra_lazy(g, 0)[0] == want and dijkstra_indexed(g, 0)[0] == want
    stairs = staircase_graph(12)
    assert dijkstra_lazy(stairs, 0)[0] == dijkstra_indexed(stairs, 0)[0] == bellman_ford(stairs, 0)
    print("both variants match Bellman-Ford on 300 random graphs and on a 12-node staircase")

    cases = [
        ("sparse random, 2000 nodes, 10000 edges", random_graph(2000, 10000, 1)),
        ("dense random, 300 nodes, 45000 edges", random_graph(300, 45000, 2)),
        ("staircase, 300 nodes", staircase_graph(300)),
    ]
    for name, g in cases:
        edges = sum(len(x) for x in g.values())
        d1, lazy = dijkstra_lazy(g, 0)
        d2, indexed = dijkstra_indexed(g, 0)
        assert d1 == d2
        print(f"\n{name} ({edges} edges)")
        print("  lazy duplicates:", lazy)
        print("  indexed heap:   ", indexed)

Run, it prints:

lazy duplicates on a 4-node graph:
  pop (0, node 0)
  push (4, node 1); heap now [(4, 1)]
  push (1, node 2); heap now [(1, 2), (4, 1)]
  pop (1, node 2)
  push (3, node 1); heap now [(3, 1), (4, 1)]
  push (6, node 3); heap now [(3, 1), (4, 1), (6, 3)]
  pop (3, node 1)
  push (4, node 3); heap now [(4, 1), (4, 3), (6, 3)]
  pop (4, node 1)  -> stale, skipped
  pop (4, node 3)
  pop (6, node 3)  -> stale, skipped
  ({0: 0, 1: 3, 2: 1, 3: 4}, {'pushes': 6, 'stale pops': 2, 'peak size': 3, 'comparisons': 6})
indexed heap on the same graph:
  pop (0, node 0)
  push (4, node 1); heap now [(4, 1)]
  push (1, node 2); heap now [(1, 2), (4, 1)]
  pop (1, node 2)
  decrease node 1 to 3; heap now [(3, 1)]
  push (6, node 3); heap now [(3, 1), (6, 3)]
  pop (3, node 1)
  decrease node 3 to 4; heap now [(4, 3)]
  pop (4, node 3)
  ({0: 0, 1: 3, 2: 1, 3: 4}, {'pushes': 4, 'decrease-keys': 2, 'peak size': 2, 'comparisons': 2})

both variants match Bellman-Ford on 300 random graphs and on a 12-node staircase

sparse random, 2000 nodes, 10000 edges (10000 edges)
  lazy duplicates: {'pushes': 2759, 'stale pops': 771, 'peak size': 1328, 'comparisons': 51174}
  indexed heap:    {'pushes': 1988, 'decrease-keys': 771, 'peak size': 969, 'comparisons': 35954}

dense random, 300 nodes, 45000 edges (45000 edges)
  lazy duplicates: {'pushes': 947, 'stale pops': 647, 'peak size': 868, 'comparisons': 16293}
  indexed heap:    {'pushes': 300, 'decrease-keys': 647, 'peak size': 289, 'comparisons': 5102}

staircase, 300 nodes (44850 edges)
  lazy duplicates: {'pushes': 44851, 'stale pops': 44551, 'peak size': 44552, 'comparisons': 1314931}
  indexed heap:    {'pushes': 300, 'decrease-keys': 44551, 'peak size': 299, 'comparisons': 48000}

Reading the trace

The 4-node graph has five edges: 0 to 1 (weight 4), 0 to 2 (weight 1), 1 to 3 (weight 1), 2 to 1 (weight 2) and 2 to 3 (weight 5); node 3 has no outgoing edges. By hand: from node 0, node 1 costs 4 and node 2 costs 1. Popping node 2 (distance 1) relaxes 2 to 1 giving 1 + 2 = 3, better than 4, and relaxes 2 to 3 giving 1 + 5 = 6. Popping node 1 (distance 3) relaxes 1 to 3 giving 3 + 1 = 4, better than 6.

On the small graph both variants find the distances {0: 0, 1: 3, 2: 1, 3: 4}. Node 1 is first reached at 4 through the direct edge, then improved to 3 through node 2. The indexed heap changes its entry in place (decrease node 1 to 3), so the heap never holds two entries for node 1 and has 2 entries at its peak. The lazy heap pushes (3, node 1) next to the old (4, node 1), and later pops (4, node 1) and skips it. The same happens for node 3, whose entry (6, node 3) goes stale when (4, node 3) is pushed. Lazy: 6 pushes, 2 stale pops, peak 3. Indexed: 4 pushes, 2 decrease-keys, peak 2.

Trade-offs

Measured on the graphs above (heap comparisons, lazy against indexed):

• Sparse random (2000 nodes, 10000 edges): 51174 against 35954, about 1.4 times, peak sizes 1328 against 969.
• Dense random (300 nodes, 45000 edges): 16293 against 5102, about 3.2 times, peak sizes 868 against 289.
• Staircase (300 nodes, 44850 edges): 1314931 against 48000, about 27 times, with a lazy peak of 44552 entries against 299.

The staircase is the worst case for duplicates: every edge improves a distance, so the lazy heap grows to E entries. On ordinary graphs the number of improvements is small, so the duplicates cost a modest constant factor.

Recommendation

For typical sparse graphs (E close to V times a small constant), use duplicates with skipping on the library heap. It is a few lines, it cannot corrupt a position map, and heapq runs in C, so its raw speed usually outweighs the extra comparisons that the counts above show (the counts compare algorithms, not wall-clock time). Use the indexed heap when the graph is dense or when edge count makes O(E) heap memory unacceptable, when you need a worst-case bound that does not depend on how many improvements occur, or when the queue is long-lived and also serves other priority updates by id. A binary heap's decrease-key is O(log V). A Fibonacci heap (an advanced textbook structure that you will rarely be asked to build) does it in O(1) amortized, meaning O(1) when averaged over a long sequence of operations, giving O(E + V log V) overall, but its large constants mean it is rarely used in practice.

Edge cases and common mistakes

• Negative edge weights break Dijkstra itself, not only the queue.
• Skipping stale entries must happen before the neighbours are relaxed. Relaxing from a stale entry scans all of that node's edges a second time for nothing.
• Forgetting to update pos inside _swap (or in pop for the last entry moved to the root) leaves wrong positions and decrease-key changes the wrong entry.
• Calling decrease_key on a node that was already popped: the position map no longer contains it, so check membership first.
• Ties: with [priority, item] entries Python compares the items when priorities are equal, which fails for non-comparable nodes. Add a counter as a tie-breaker in that case.

Running the code

python decrease_key.py

Python 3.12, standard library only (random). No install or config needed.

The idea

Two heaps split the numbers into a lower half and an upper half. lowers is a max-heap (a heap whose top is its largest item) and highers is a min-heap (top is its smallest item). Two rules must hold after each insertion:

1. Ordering: every number in lowers is less than or equal to every number in highers, i.e. the top of lowers <= the top of highers.
2. Size: the halves are equal in size, or lowers has exactly one more. (This convention makes the median the top of lowers when the count is odd.)

Start: lowers = [5, 3, 2] (top 5), highers = [6, 8, 9] (top 6). Both rules hold: 5 <= 6, and the sizes are 3 and 3. The median of 2, 3, 5, 6, 8, 9 is (5 + 6) / 2 = 5.5.

Steps for inserting 7

Step 1: decide which half 7 belongs to. Compare with the top of lowers. If 7 were less than or equal to 5 it would go in lowers; it is larger, so it goes in highers. Putting it there cannot break the ordering rule, because 7 is at least as big as everything in lowers. highers becomes {6, 7, 8, 9} with top 6.

Step 2: check the sizes. lowers has 3 items and highers has 4. The size rule says highers may never have more items than lowers, so this breaks it.

Step 3: rebalance by moving one item. Pop the top of highers (6, the smallest of the upper half) and push it into lowers. Moving the smallest of the upper half is the only safe choice: it is at least as big as everything in lowers and no bigger than anything left in highers, so ordering still holds. The tops become 6 and 7, and 6 <= 7.

Step 4: read the result. Sizes are now 4 and 3, which is allowed. With 7 numbers (2, 3, 5, 6, 7, 8, 9) the median is the top of lowers, 6.

The script below runs these four steps, numbered the same way, on real heapq heaps (the max-heap is stored negated because heapq is a min-heap) and prints the state after each one:

import heapq

lowers = [-v for v in (5, 3, 2)]      # max-heap, stored negated
highers = [6, 8, 9]                   # min-heap
heapq.heapify(lowers)
heapq.heapify(highers)


def show(label):
    print(f"{label:<34} lowers(max-heap)={sorted((-v for v in lowers), reverse=True)} "
          f"top={-lowers[0]}  highers(min-heap)={sorted(highers)} top={highers[0]}")


def median():
    return -lowers[0] if len(lowers) > len(highers) else (-lowers[0] + highers[0]) / 2


show("start")
x = 7
if x <= -lowers[0]:
    heapq.heappush(lowers, -x)
    print("step 1: 7 <= 5? yes, goes to lowers")
else:
    heapq.heappush(highers, x)
    print("step 1: 7 <= 5? no, goes to highers")
show("after step 1")
print(f"step 2: sizes are lowers={len(lowers)} highers={len(highers)}")
if len(highers) > len(lowers):
    moved = heapq.heappop(highers)
    heapq.heappush(lowers, -moved)
    print(f"step 3: highers is larger, move its top {moved} to lowers")
show("after step 3")
print(f"step 4: sizes are lowers={len(lowers)} highers={len(highers)}, median={median()}")

start                              lowers(max-heap)=[5, 3, 2] top=5  highers(min-heap)=[6, 8, 9] top=6
step 1: 7 <= 5? no, goes to highers
after step 1                       lowers(max-heap)=[5, 3, 2] top=5  highers(min-heap)=[6, 7, 8, 9] top=6
step 2: sizes are lowers=3 highers=4
step 3: highers is larger, move its top 6 to lowers
after step 3                       lowers(max-heap)=[6, 5, 3, 2] top=6  highers(min-heap)=[7, 8, 9] top=7
step 4: sizes are lowers=4 highers=3, median=6

Final heaps: lowers = {6, 5, 3, 2} (top 6), highers = {7, 8, 9} (top 7).

Why not other moves

• Pushing 7 straight into lowers would put 7 above 6, the smallest of highers, and break the ordering rule.
• Rebalancing by moving the top of lowers (5) into highers is the move for the opposite imbalance, and here it would make the size gap worse.
• Under a looser rule (sizes may differ by one in either direction) the 3 versus 4 split is already legal and step 3 is skipped; the median is then the top of highers, which is also 6. Pick one convention and keep it, because median() depends on it.

Complexity and mistakes

Insertion and rebalancing are a constant number of heap operations, so O(log n) per insert, with O(1) to read the median. Common mistakes: forgetting that a move changes both heaps (pop from one and push to the other), and rebalancing before routing by value.

Running the code

python3 insert_trace.py

Python 3.12, standard library only.