Amazon — Kth Largest Element in a Stream (Min-Heap)

Problem Statement#

Design a class KthLargest that finds the kth largest element in a stream of integers. The kth largest is the kth largest in sorted order, not the kth distinct element.

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(-3), add(-2), add(-4), add(0), add(4)
Output: [-3,-2,-2,0,4]

Why This Problem Matters#

Kth Largest Element in a Stream (LeetCode 703) is an Amazon interview staple that directly models real-world streaming analytics. Amazon's recommendation engine, order ranking, and real-time bidding systems all need to track top-K metrics over a continuous stream without storing all historical data. This problem tests whether you understand that a min-heap of size k elegantly solves this in O(log k) per insertion.

The naive approach — sort all elements for every add call — costs O(N log N) per call, which is completely impractical for a stream. The heap approach maintains exactly k elements: the k largest seen so far. The minimum of those k elements is always the kth largest.

Google, Microsoft, and Uber also ask this in variants: "find the kth smallest," "maintain a sliding window percentile," or "track the top-k most frequent items." All use the same heap-size-k invariant.

The Core Insight#

Maintain a min-heap of exactly k elements. The invariant: the heap always contains the k largest elements seen so far, and the heap minimum (top of min-heap) is the kth largest.

When a new element arrives: push it onto the heap. If the heap size exceeds k, pop the minimum. After this, the heap top is always the kth largest. The heap size never exceeds k, bounding both space and per-operation time.

Visual Dry Run#

k=3, initial nums=[4,5,8,2]

Solution (Optimal)#

import heapq
 
class KthLargest:
    def __init__(self, k: int, nums: list):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)
 
    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.data = []; }
    size() { return this.data.length; }
    peek() { return this.data[0]; }
    push(val) {
        this.data.push(val);
        let i = this.data.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.data[p] <= this.data[i]) break;
            [this.data[p], this.data[i]] = [this.data[i], this.data[p]];
            i = p;
        }
    }
    pop() {
        const top = this.data[0];
        const last = this.data.pop();
        if (this.data.length > 0) {
            this.data[0] = last;
            let i = 0;
            while (true) {
                let min = i, l = 2*i+1, r = 2*i+2;
                if (l < this.data.length && this.data[l] < this.data[min]) min = l;
                if (r < this.data.length && this.data[r] < this.data[min]) min = r;
                if (min === i) break;
                [this.data[min], this.data[i]] = [this.data[i], this.data[min]];
                i = min;
            }
        }
        return top;
    }
}
 
class KthLargest {
    constructor(k, nums) {
        this.k = k;
        this.heap = new MinHeap();
        for (const num of nums) this.add(num);
    }
    add(val) {
        this.heap.push(val);
        if (this.heap.size() > this.k) this.heap.pop();
        return this.heap.peek();
    }
}

Time: O(log k) per add call — heap operations on a heap of size k Space: O(k) — heap stores exactly k elements

Common Mistakes#

• Using a max-heap of size k — gives the kth smallest, not kth largest
• Not populating the heap through add in __init__ — misses the initialization logic
• Returning heap[0] before k elements exist — undefined behavior on small initial lists
• Forgetting that k can exceed the initial array length — heap starts smaller than k
• Using sorted insertion instead of a heap — O(N) per insert vs O(log k)

Interview Tips#

• Say it clearly: min-heap of size k gives kth largest (the heap minimum is the answer)
• Python heapq is a min-heap by default — no negation needed
• JavaScript needs a manual heap implementation — mention this upfront and code it quickly
• Test your edge case: what if fewer than k elements have been added? (heap still correct)
• Bring up the follow-up immediately: kth smallest uses a max-heap of size k

Follow-up Questions#

• How do you find the kth smallest? — Use a max-heap of size k; the max is the kth smallest
• What if k changes dynamically? — Rebuild heap; no incremental update possible
• How do you handle deletions from the stream? — Lazy deletion with a secondary "removed" set
• What if you need all k elements, not just the kth? — Return all heap elements sorted
• How do you solve this for k=1? — Running maximum; just track a single variable

Key Takeaways#

• A min-heap of exactly k elements keeps the k largest values; its minimum is the kth largest
• Each add operation costs O(log k) — independent of total stream length seen so far
• Space is O(k) — the heap never grows beyond k elements regardless of stream volume
• Python's heapq is a min-heap natively; JavaScript interviews require a manual implementation
• Amazon tests this to verify streaming algorithm thinking and heap data structure mastery
• The min-heap-of-size-k pattern generalizes to any "maintain top-k" or "kth order statistic" problem
• Initializing via repeated add calls costs O(N log k); using heapify then trimming is O(N)