Top-K — Heap Is Not Always the Answer

Finding the K-th largest element in an integer array (LeetCode 215) is the canonical top-K problem. My reflex for this shape of problem was always a heap. Re-solving the same problem with Quick Select showed that the reflex is only half right. The two tools reach the same answer at different costs: with streaming data a heap is effectively the only option, while with the whole array already in memory Quick Select is faster on average.

Full sorting works too. With N around 10⁵, an O(N log N) sort passes. But fixing the order of all N elements to extract a single K-th value is more work than the problem asks for.

Size-K Min-Heap

The heap solution hinges on a min-heap, not a max-heap. Keep only the K largest values seen so far in a min-heap of size K; the root heap[0] is then the smallest of those K, which is exactly the current candidate for the K-th largest. The first K values go in unconditionally; after that, each new value is compared against the root. When a new value exceeds the root, drop the root and insert the new value. After the pass, the root is the answer.

Building a max-heap over all N elements and popping K times also works, but the space grows to O(N). The size-K approach stays at O(K), and the advantage widens as K gets much smaller than N.

flowchart TD
    A["new value"] --> B{"heap size < K"}
    B -->|yes| C["heappush"]
    B -->|no| D{"value > heap[0]"}
    D -->|yes| E["heapreplace"]
    D -->|no| F["discard"]

import heapq


def find_kth_largest(nums, k):
    top_k = []
    for value in nums:
        if len(top_k) < k:
            heapq.heappush(top_k, value)
        elif value > top_k[0]:
            heapq.heapreplace(top_k, value)
    return top_k[0]

heapreplace performs the pop and the push in a single reheapify; calling heappop then heappush separately reorders the heap twice. The strict > comparison is deliberate: a value equal to the root would not change the candidate, so it is skipped.

The strength of this structure is streaming. Each arriving element goes through the same check, and N never needs to be known in advance. The same code keeps working for a live leaderboard or a log stream, where elements keep arriving and the end is never known, and reading top_k[0] at any moment gives the K-th largest so far. The cost is O(log K) per element, O(N log K) overall.

Quick Select

Quick Select starts from a different observation: the K-th largest value is the element at index N - K in ascending order (0-indexed). In [3, 2, 1, 5, 6, 4] the 2nd largest is 5, which lands at index 4 after sorting. If that one position can be fixed without sorting everything, the sorting cost disappears, and Quick Sort’s partition can fix it.

partition picks a pivot, moves everything smaller than it to the left, places the pivot at its final sorted position, and returns that position. If the returned position equals the target index, that value is the answer. Otherwise, partition again only in the half that contains the target. Where Quick Sort recurses into both halves, Quick Select descends into one, and that difference sets the average complexity: with the range roughly halving each time, N + N/2 + N/4 + … stays under 2N. Average O(N).

flowchart TD
    A["partition(lo, hi) → p"] --> B{"p == target"}
    B -->|yes| C["answer nums[p]"]
    B -->|"p < target"| D["partition the right range only"]
    B -->|"p > target"| E["partition the left range only"]
    D --> A
    E --> A

import random


def find_kth_largest(nums, k):
    target = len(nums) - k
    lo, hi = 0, len(nums) - 1
    while True:
        pivot_pos = random.randint(lo, hi)
        nums[pivot_pos], nums[hi] = nums[hi], nums[pivot_pos]
        pivot = nums[hi]
        boundary = lo
        for i in range(lo, hi):
            if nums[i] < pivot:
                nums[boundary], nums[i] = nums[i], nums[boundary]
                boundary += 1
        nums[boundary], nums[hi] = nums[hi], nums[boundary]
        if boundary == target:
            return nums[boundary]
        if boundary < target:
            lo = boundary + 1
        else:
            hi = boundary - 1

target is the sorted-order index of the K-th largest value, and boundary is the next slot of the region holding values smaller than the pivot. The line that picks the pivot at random is what prevents the worst case. On already-sorted input, always taking the last element as pivot shrinks the range by one position per round, giving O(N²). Randomization keeps the O(N) average for any input.

One caveat: this is in-place. The input array gets shuffled as a side effect. Preserving the original order means starting from a copy, and the space advantage disappears with it.

Selection Criteria

When the heap fits:

• Data arrives as a stream, or N is unknown in advance
• The K-th value is queried repeatedly as data comes in
• The original array must stay intact and space should stay at O(K)

When Quick Select fits:

• The whole array is already in memory
• The K-th value is needed once
• Mutating the array in place is acceptable

In a live-coding setting, I would pass the problem with the short, predictable heap solution first, then mention Quick Select’s O(N) average and pivot randomization if time remains.

The problem statement stays the same — K-th largest — but the shape of the data picks the tool. For a stream, maintain a heap of size K; for an array already in memory, converge on the target index with partition. What this re-solve left me with is the habit of checking the data’s shape before defaulting to a heap.