Kth Largest Element: Heap vs Quickselect

Pattern: Heap (Top-K and Greedy Merging) · Difficulty: Medium · Asked at: Meta, Amazon, LinkedIn, Spotify

Classic version: LeetCode 215 · LeetCode 703

Problem

1. kth_largest(nums, k): return the k-th largest value in an unsorted list (the k-th in sorted-descending order, duplicates counted).
2. KthLargestStream(k, initial): a class whose add(x) inserts a value and returns the current k-th largest of everything seen so far. Assume at least k values exist after each add.

Examples

kth_largest([3, 2, 1, 5, 6, 4], 2)          → 5
kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4) → 4

s = KthLargestStream(3, [4, 5, 8, 2])
s.add(3) → 4    s.add(5) → 5    s.add(10) → 5    s.add(9) → 8

Starter code

def kth_largest(nums: list[int], k: int) -> int:
    pass


class KthLargestStream:
    def __init__(self, k: int, initial: list[int]):
        pass

    def add(self, x: int) -> int:
        pass

Hints

Where this shows up in data engineering

“Top 10 customers by spend”, “p99 latency”, “largest 100 files to compact first”. In a stream you keep a bounded heap per key (that’s what approx_top_k and leaderboard services do internally). In batch, Spark’s takeOrdered keeps a size-k heap per partition and merges them, which is why it beats a full orderBy().limit() for small k.

Solution

import heapq
import random


def kth_largest(nums, k):
    # quickselect: find the element that would sit at index n-k in ascending order
    target = len(nums) - k
    a = list(nums)
    lo, hi = 0, len(a) - 1
    while True:
        pivot = a[random.randint(lo, hi)]
        # three-way partition handles duplicates without degrading
        lt = [x for x in a[lo:hi + 1] if x < pivot]
        eq = [x for x in a[lo:hi + 1] if x == pivot]
        gt = [x for x in a[lo:hi + 1] if x > pivot]
        a[lo:hi + 1] = lt + eq + gt
        if target < lo + len(lt):
            hi = lo + len(lt) - 1
        elif target < lo + len(lt) + len(eq):
            return pivot
        else:
            lo = lo + len(lt) + len(eq)


class KthLargestStream:
    def __init__(self, k, initial):
        self.k = k
        self.heap = []                     # min-heap holding the k largest values
        for x in initial:
            self.add(x)

    def add(self, x):
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, x)
        elif x > self.heap[0]:
            heapq.heapreplace(self.heap, x)  # pop root + push in one O(log k) step
        return self.heap[0]

Tests

Your solution should pass these:

assert kth_largest([3, 2, 1, 5, 6, 4], 2) == 5
assert kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4) == 4
assert kth_largest([1], 1) == 1
assert kth_largest([7, 7, 7, 7], 2) == 7
assert kth_largest(list(range(1000)), 1) == 999
s = KthLargestStream(3, [4, 5, 8, 2])
assert [s.add(3), s.add(5), s.add(10), s.add(9), s.add(4)] == [4, 5, 5, 8, 8]
t = KthLargestStream(1, [])
assert [t.add(-3), t.add(-2), t.add(-4)] == [-3, -2, -2]

Explanation

Heap (stream or array): a min-heap of size k keeps exactly the k largest seen so far; the smallest of them, the root, is the answer. O(n log k) time, O(k) memory, and it works on unbounded streams. In Python heapq.nlargest(k, nums)[-1] does this.

Quickselect (array only): partition around a random pivot; only recurse into the side containing index n-k. Expected O(n), worst case O(n²) (mitigated by random pivots; median-of-medians guarantees O(n) but is rarely asked). The three-way partition prevents quadratic behaviour on many duplicates.

Sorting: O(n log n), fine as a first answer; say why you’d improve it.

Choosing: stream or huge n with small k → heap. One-off in-memory array → quickselect. Need the whole ranking anyway → sort.

Follow-up questions