The Two Heaps Pattern

The two heaps pattern keeps one max heap and one min heap and maintains an invariant that divides the data between them. The classic form for a running median keeps the lower half in the max heap and the upper half in the min heap, with the max heap holding either the same number of elements as the min heap or exactly one more.

That invariant gives you the median directly: the root of the max heap is the largest of the lower half, and if the two heaps are the same size the median is the average of that root and the root of the min heap. Every insertion and every removal touches only the heap roots, so both operations are logarithmic.

The same pattern solves other problems. A median of a sliding window uses it with an extra structure to expire old elements. A kth smallest or kth largest in a stream uses one heap with a size cap. Two heaps can also maintain the smallest and largest values of a collection at the same time, or split values by some threshold into a below group and an above group.

The trade-off to name is that two heaps hold the data twice over in terms of bookkeeping, so you pay roughly double the memory of a single heap plus the effort of keeping the invariant after every operation. The invariant itself is the hard part: every insertion and removal has to check which side the element belongs on and then rebalance by moving one element across. That rebalancing step is where most bugs in this pattern live.

Keep a max heap for the lower half and a min heap for the upper half, and rebalance so the lower half is never smaller than the upper half and never more than one element larger.

import heapq


class MedianFinder:
    def __init__(self):
        self.lower = []
        self.upper = []

    def add(self, value):
        if not self.lower or value <= -self.lower[0]:
            heapq.heappush(self.lower, -value)
        else:
            heapq.heappush(self.upper, value)
        if len(self.lower) > len(self.upper) + 1:
            heapq.heappush(self.upper, -heapq.heappop(self.lower))
        else:
            if len(self.upper) > len(self.lower):
                heapq.heappush(self.lower, -heapq.heappop(self.upper))

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


finder = MedianFinder()
for value in [5, 15, 1, 3]:
    finder.add(value)
print(finder.median())

The output is 4.0. After the four insertions the lower heap holds 1 and 3 and the upper heap holds 5 and 15, so with an even number of values the median is the average of the two middle ones, which is four.

The mechanism is that the lower heap is negated so that a min heap behaves like a max heap, which lets you read its largest element from the root. The insertion rule puts the new value on the correct side by comparing it with the root of the lower heap, and the rebalance step moves one element across the boundary whenever the sizes drift apart.

The trade-off to name is that the median of an even number of values is the average of the two middle ones, so the query has to look at both roots, and you must decide what to return before any value has been added, since the median of an empty stream is not defined. The other detail is that the negation trick only works for numbers, so for arbitrary objects you would need a wrapper class with an inverted comparison.

Count the occurrences, then push each distinct value together with its negative count onto a heap so the most frequent one comes out first.

from collections import Counter
import heapq


def top_k_frequent(values, k):
    counts = Counter(values)
    heap = [(-count, value) for value, count in counts.items()]
    heapq.heapify(heap)
    result = []
    for _ in range(k):
        if not heap:
            break
        count, value = heapq.heappop(heap)
        result.append(value)
    return result


print(top_k_frequent([1, 1, 1, 2, 2, 3], 2))

The output is 1 and 2. The value 1 occurs three times, 2 occurs twice and 3 once, so the two most frequent are 1 and 2.

The mechanism is the negated count as the first element of the tuple. A min heap surfaces the smallest tuple, and the smallest negated count is the largest count, so the most frequent item comes out first. The value is carried as the second element and also breaks ties, so two items with the same count are ordered by their value.

The trade-off to name is that negating the counts is a trick that only works for integers, and it costs a list comprehension over all the distinct values, so you build a heap of size equal to the number of distinct items even though you only want k of them. If k is much smaller than the number of distinct items, the better approach is to keep a heap of size k and replace its root as you go, which is the version shown in the heap group earlier. Both are correct; they differ in memory and in how many heap operations they perform.