Kth Largest Element in a Stream (LeetCode 703)

The problem

Design a class KthLargest(k, nums). It starts with the numbers in nums. Each call to add(val) adds one more number to the stream and returns the kth largest number seen so far (counting duplicates).

Input:  KthLargest(3, [4, 5, 8, 2])
        add(3), add(5), add(10), add(9), add(4)

Output: 4, 5, 5, 8, 8

After add(3) the stream is 2, 3, 4, 5, 8; the three largest are 8, 5, 4, so the third largest is 4.

How to think about it

Start with brute force. Keep every number in a list. On each add, sort it and read position k from the end. That’s O(n log n) per call, and the stream can grow without limit. Keeping a sorted list and inserting with binary search is better, but the insert still shifts O(n) items.

The key insight: you only care about the top k. Any number smaller than the current kth largest can never matter again, since new numbers only push it further down. So keep just the k largest numbers. Among those, the kth largest is the smallest, so store them in a min-heap: its root is the answer, readable in O(1).

Each add is one push and maybe one pop. Push the new value. If the heap now has k + 1 items, pop the smallest; it has just fallen out of the top k. That’s the lesson’s signature trick: a min-heap of size k keeps the k largest. It feels backwards the first time, which is why it’s worth making a reflex.

Which pattern is this? Top-K with a heap. The tell is “kth largest” plus data that keeps arriving, so you can’t sort once and be done.

Approach

1. In the constructor, copy nums, heapify it, and pop until only k items remain.
2. In add(val), push val onto the heap.
3. If the heap has more than k items, pop the smallest.
4. Return the root, h[0].

Solution

import heapq
class KthLargest:
    def __init__(self, k, nums):
        self.k = k; self.h = nums[:]; heapq.heapify(self.h)
        while len(self.h) > k: heapq.heappop(self.h)        # keep only the k largest
    def add(self, val):
        heapq.heappush(self.h, val)
        if len(self.h) > self.k: heapq.heappop(self.h)      # smallest drops out of the top k
        return self.h[0]                                     # root = kth largest
# KthLargest(3, [4,5,8,2]); add 3,5,10,9,4 -> 4,5,5,8,8

Dry run

The constructor heapifies [4, 5, 8, 2] and pops 2, leaving the top three: [4, 5, 8].

Notice add(3) and add(4): the new value is pushed and immediately popped, because it’s too small to join the top three.

Complexity

• Time: O(log k) per add. The heap never holds more than k + 1 items. The constructor is O(n + (n - k) log n) for heapify plus the pops.
• Space: O(k). Only the k largest values are stored.

Edge cases and pitfalls

• Fewer than k numbers at the start. With k = 2, nums = [0] the heap has one item, and the first add must not pop. The > self.k check handles this; LeetCode guarantees there are at least k numbers whenever a result is read.
• Empty nums works the same way: the heap fills up through add.
• Duplicates count separately. Two 5s take up two places in the top k, as the second add shows.
• Max-heap of everything is the common wrong instinct. It needs O(n) space and O(k log n) to read the kth item.
• Copy the input (nums[:]) so heapifying doesn’t reorder the caller’s list.

Related problems

• Kth Largest Element in an Array (LeetCode 215): the same size-k heap over a fixed array.
• Find Median from Data Stream (LeetCode 295): another streaming statistic, using two heaps.
• Top K Frequent Elements (LeetCode 347): the size-k heap keyed by frequency.