Kth Largest Element in a Stream — LeetCode 703 Heap Pattern

Sanjeev SharmaSanjeev Sharma
5 min read

Advertisement

Problem Statement

Design a class that returns the Kth largest element in a stream after each insertion. Initialize with a stream prefix and an integer K, then add new values one by one.

Constraints:

  • 1 <= k <= 10^4
  • 0 <= nums.length <= 10^4
  • -10^4 <= nums[i] <= 10^4
  • At most 10^4 calls to add
Input:  k = 3, nums = [4, 5, 8, 2]
        add(3) add(5) add(10) add(9) add(4)
Output: 4, 5, 5, 8, 8
Input:  k = 1, nums = []
        add(-1) add(2) add(0)
Output: -1, 2, 2

Why This Problem Matters

LeetCode 703 is one of the most common phone-screen warmups at Amazon, Meta, and Microsoft. It crystallizes the fixed-size min-heap pattern, which scales to harder problems like Top K Frequent Elements and K Closest Points to Origin. Recognizing this pattern in 30 seconds signals heap fluency to interviewers.

The problem also tests API design: you must build a class that maintains state across calls, not a one-shot function. That's exactly what real systems (leaderboards, recommendation engines, anomaly detectors) need.

Keywords: "Kth largest interview", "streaming top K", "priority queue FAANG", "min-heap size k".

The Core Insight

The Kth largest in a stream of N values is the smallest of the top K. Maintain a min-heap of capacity K. After each insert, if the heap exceeds K, pop the smallest. The root is always the answer in O(1).

Why min-heap and not max-heap? A max-heap would give you the largest each time, but you would have to scan deeper to find the Kth. A min-heap of size K naturally evicts everything smaller than the top K seen so far.

Visual Dry Run

StepAddHeap (min, size <= 3)Kth Largest
init4, 5, 8, 24, 5, 84
134, 5, 8 (3 popped)4
255, 5, 85
3105, 8, 105
498, 9, 108
548, 9, 108

Solution (Optimal)

import heapq
 
class KthLargest:
    def __init__(self, k: int, nums):
        self.k = k
        self.heap = []
        for v in nums:
            self.add(v)
 
    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]
class MinHeap {
    constructor() { this.h = []; }
    push(v) { this.h.push(v); this._up(this.h.length - 1); }
    pop() {
        const top = this.h[0], last = this.h.pop();
        if (this.h.length) { this.h[0] = last; this._down(0); }
        return top;
    }
    peek() { return this.h[0]; }
    size() { return this.h.length; }
    _up(i) {
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.h[p] <= this.h[i]) break;
            [this.h[p], this.h[i]] = [this.h[i], this.h[p]];
            i = p;
        }
    }
    _down(i) {
        const n = this.h.length;
        while (true) {
            const l = 2 * i + 1, r = 2 * i + 2;
            let s = i;
            if (l < n && this.h[l] < this.h[s]) s = l;
            if (r < n && this.h[r] < this.h[s]) s = r;
            if (s === i) break;
            [this.h[s], this.h[i]] = [this.h[i], this.h[s]];
            i = s;
        }
    }
}
 
var KthLargest = function(k, nums) {
    this.k = k;
    this.heap = new MinHeap();
    for (const v of nums) this.add(v);
};
 
KthLargest.prototype.add = function(val) {
    this.heap.push(val);
    if (this.heap.size() > this.k) this.heap.pop();
    return this.heap.peek();
};

Time: O(log k) per add, O(n log k) for initialization with n values. Space: O(k) — heap holds at most K elements regardless of stream length.

Common Mistakes

  • Using a max-heap of size N — that becomes O(n) memory and O(log n) time per add unnecessarily.
  • Sorting on every add — O(n log n) per call, fails the constraints.
  • Forgetting to call add for the initial nums in the constructor.
  • Returning the wrong index after pop instead of the heap root.
  • Not handling the case where initial nums has fewer than K elements (the heap simply has not filled yet).

Interview Tips

  • Say out loud: "I want O(log k) per add and O(k) space — that's a min-heap of size K."
  • Reuse the add method inside the constructor to avoid duplicating logic.
  • Mention that this generalizes to any "top K streaming" problem.
  • Note that ties do not matter — the Kth largest is well-defined even with duplicates.

Follow-up Questions

  • What if K can change at runtime? Store K mutably and shrink or grow the heap on change.
  • What if you need both Kth largest AND Kth smallest? Maintain two heaps.
  • What if values can be removed? Use a hash with lazy deletion.
  • What if the stream is distributed across machines? Each shard maintains a top-K heap; merge with a final K-way merge.

Key Takeaways

  • LeetCode 703 is the canonical "top K streaming" problem.
  • A min-heap of size K gives O(log k) per add and O(k) space.
  • The root of the size-K min-heap is always the Kth largest element seen so far.
  • This pattern beats sorting on every add and beats keeping all elements in a max-heap.
  • Reuse the add method inside the constructor to keep code DRY.
  • Generalizes to Top K Frequent and K Closest Points.
  • Python heapq makes this a 5-line solution; JavaScript needs a heap class.

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro

Related reading