Kth Largest Element in a Stream

1. Problem & Core Objective

Design a class that reports the k-th largest element seen so far, as values stream in.

KthLargest(int k, int[] nums)   // seed with an initial array
int add(int val)                // add a value, return the kth largest so far

KthLargest kth = new KthLargest(3, [4, 5, 8, 2]);
kth.add(3);   → 4        // sorted: 8,5,4,3,2 → 3rd largest is 4
kth.add(5);   → 5        // 8,5,5,4,3,2      → 5
kth.add(10);  → 5
kth.add(9);   → 8
kth.add(4);   → 8

Constraints: 1 <= k <= 10^4 · 0 <= nums.length <= 10^4 · up to 10^4 calls to add · it is guaranteed there are always at least k elements when add is called

What's actually being tested: the counterintuitive pairing — to track the k LARGEST, you use a MIN-heap. If that inversion doesn't feel obvious, this is the question to fix it on, because questions 3 and 4 rely on the same move.

2. First-Principles Thought Process

What "streaming" rules out

Values arrive one at a time and the answer is needed after each. So any approach that re-derives the answer from scratch pays that cost on every call.

Sorting is the obvious example: O(n log n) per add, with n growing. That's not a constant-factor problem, it's the wrong shape.

The observation that shrinks the problem

To report the k-th largest, you only ever need the top k values. Everything below them is irrelevant — it can never become the k-th largest, because k values already beat it and they're all still present.

So: keep exactly k elements, discard the rest permanently. Memory drops from O(n) to O(k).

Which k elements, and which one is the answer?

Keep the k largest. Then the k-th largest is the smallest of those k — and that's the one you need constant access to.

A structure whose root is the minimum is a min-heap.

Why a max-heap is exactly wrong

A max-heap puts the largest on top. But the largest is the element you most want to keep, and it's the one you'd never evict. The element you actually need — both to report and to discard — is the smallest of the kept set, which a max-heap buries at a leaf where finding it is O(k).

Min-heap for the k largest; max-heap for the k smallest. The heap is inverted relative to the question, and that inversion is the whole insight.

The maintenance rule

minHeap.offer(val);
if (minHeap.size() > k) minHeap.poll();   // drop the smallest
return minHeap.peek();                     // the kth largest

Offer unconditionally, then trim. If the new value belonged in the top k, something else gets pushed out; if it didn't, it is itself immediately evicted. Either way the invariant "the heap holds the k largest seen" is restored in one line.

3. Solution Paths

Approach 1 — Keep everything, sort on each query (brute force)

class KthLargest {
    private final List<Integer> all = new ArrayList<>();
    private final int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int n : nums) all.add(n);
    }

    public int add(int val) {
        all.add(val);
        Collections.sort(all, Collections.reverseOrder());
        return all.get(k - 1);
    }
}

• Time O(n log n) per add · Space O(n)

Counter-questions on this approach

⭐ "What's the total cost across all calls?"

With 10^4 adds against a collection growing to 2 × 10^4, each sort is about 2 × 10^4 × 15 ≈ 3 × 10^5 comparisons, so roughly 3 × 10^9 overall. Far too slow.

The deeper problem is that it re-derives the full order every time to read one element. Sorting produces complete information; the question asks for one rank.

⭐ "You could insert into a sorted list instead of re-sorting. Does that fix it?"

It helps but doesn't solve it. Binary search finds the insertion point in O(log n), but an ArrayList insert shifts elements — O(n) per add, so 2 × 10^8 overall. Borderline.

And it still stores everything. The real saving is noticing that elements outside the top k are permanently irrelevant and can be thrown away, which no sorted-list variant exploits.

"Is storing all n ever necessary?"

Only if the query could change — "give me the j-th largest for arbitrary j" would need everything. With k fixed at construction, the bottom n − k elements can never be the answer.

Approach 2 — Min-heap of size k (optimal)

class KthLargest {
    private final PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    private final int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int n : nums) add(n);         // reuse the same maintenance rule
    }

    public int add(int val) {
        minHeap.offer(val);
        if (minHeap.size() > k) minHeap.poll();
        return minHeap.peek();
    }
}

Trace — k = 3, seeded with [4,5,8,2], then add(3):

Note that 2 and 3 are evicted the instant they arrive — they never displace anything.

• Time O(log k) per add, O(n log k) to seed · Space O(k)

Counter-questions on this approach

⭐ "Why a min-heap when the question asks for the largest?"

Because of what I need constant access to. I keep the k largest values, and the answer — the k-th largest — is the smallest of those. A min-heap puts exactly that at the root, readable in O(1) and removable in O(log k).

A max-heap would surface the biggest value, which is neither the answer nor the eviction candidate. The element I need would sit somewhere among the leaves, costing O(k) to find.

The rule to remember: min-heap for the k largest, max-heap for the k smallest. The heap is always inverted relative to the wording.

⭐ "Why offer first and then check the size, rather than comparing against the root first?"

Both work. The conditional version is:

Java

if (minHeap.size() < k) minHeap.offer(val);
else if (val > minHeap.peek()) { minHeap.poll(); minHeap.offer(val); }

That avoids one heap operation when val is too small, which is a real saving if most values are. But it's three branches instead of two lines, and the offer-then-trim version restores the invariant unconditionally — harder to get wrong. I'd write the simple one and mention the optimisation.

"Why does the constructor call add rather than heapifying?"

Mostly for reuse — one maintenance rule, one place to get it right. Seeding via add is O(n log k).

If nums were large you could do better: heapify all of it in O(n) with new PriorityQueue<>(list), then poll down to size k, which is O(n + (n−k) log n). That's worse than O(n log k) when k is small, so the simple version is usually right. Worth knowing both.

"What if add is called when fewer than k elements exist?"

peek() would return the smallest of however many there are, which isn't a k-th largest at all. The constraints guarantee it can't happen, but in real code I'd return a sentinel or throw rather than a misleading number.

"Can nums be empty?"

Yes — the constraint allows length 0. The heap starts empty, and the first few add calls just fill it. The guarantee that add is only called once k elements exist is what keeps peek() meaningful.

"Does the heap hold duplicates correctly?"

Yes. PriorityQueue permits duplicates, and duplicates matter here: with [5,5,5] and k = 3, the 3rd largest is 5. A TreeSet would silently collapse them and return the wrong answer — a real trap if you reach for a sorted set instead.

Approach 3 — TreeMap as a multiset

private final TreeMap<Integer, Integer> counts = new TreeMap<>();   // value -> multiplicity
private int size;

public int add(int val) {
    counts.merge(val, 1, Integer::sum);
    size++;
    if (size > k) {
        int lowest = counts.firstKey();
        if (counts.merge(lowest, -1, Integer::sum) == 0) counts.remove(lowest);
        size--;
    }
    return counts.firstKey();
}

• Time O(log k) per add · Space O(k) distinct values

Counter-questions on this approach

⭐ "Same complexity as the heap. Why prefer the heap?"

The heap is simpler and has better constants — an array with no node allocation, versus a red-black tree with an object per entry. And this version needs explicit multiplicity bookkeeping, because a TreeMap's keys are unique while the problem's values are not.

Where the TreeMap wins is if the interface grew: it supports floorKey, ceilingKey, ordered iteration, and removal of an arbitrary value — all things a heap cannot do. For this interface none of that is needed.

"What breaks if you use a plain TreeSet instead?"

Duplicates collapse. With k = 3 and values [5,5,5], a TreeSet holds one 5 and reports the wrong rank. It's a silent wrong answer, which makes it worse than a crash — and it's the most common wrong turn on this problem.

Comparison

4. Why the Optimal Wins

Two separate savings, and it's worth naming them apart.

Bounded memory. Elements outside the top k can never become the k-th largest, so they're discarded permanently. O(n) → O(k).

Bounded work. A heap maintains just enough order to expose one element — the minimum — rather than the full ordering. O(n log n) → O(log k) per add.

Together they turn a problem that gets slower as the stream grows into one whose per-element cost is constant in the stream length.

The framing worth keeping:

To track the k largest, keep a MIN-heap of size k. Its root is both the answer and the next thing to evict.

Questions 3 and 4 are the same move on different data.

5. Java Prerequisites

PriorityQueue is a min-heap by default

PriorityQueue<Integer> minHeap = new PriorityQueue<>();
PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());

It permits duplicates and forbids nulls. Both matter here — duplicates are required for correct ranks.

O(n) heapify from a collection

new PriorityQueue<>(existingList);    // O(n), cheaper than n inserts

Comparator overflow — prefer Comparator.reverseOrder() or comparingInt over (a, b) -> b - a, which overflows for large values. See 04.

6. Interview Communication Guide

Clarifying questions: Is k fixed at construction (yes — that's what lets me discard)? Can values repeat (yes; this rules out a TreeSet)? How many add calls (10^4)? Is add guaranteed to be called only once k elements exist (yes)?

The pitch

"Because it's a stream, anything that recomputes from scratch pays that cost on every call — sorting would be O(n log n) per add with n growing.

The observation that shrinks it: to report the k-th largest I only ever need the top k values. Anything below them can never become the answer, because k larger values already exist and none of them disappear. So I can throw the rest away permanently — O(n) memory becomes O(k).

Then: which of those k do I need fast access to? The k-th largest is the smallest of the k I'm keeping — and that's also exactly the one to evict when a bigger value arrives. A structure whose root is the minimum is a min-heap.

So it's a min-heap of size k, which feels backwards given the question says 'largest'. A max-heap would surface the biggest value, which is the one I most want to keep and never want to remove — the element I actually need would be buried at a leaf.

Maintenance is two lines: offer the value, and if the size exceeds k, poll. Then peek() is the answer. A value too small to belong gets evicted immediately, so the invariant is restored either way.

O(log k) per add, O(k) space.

One thing I'd avoid: a TreeSet. Values can repeat, and a set would collapse duplicates and silently return the wrong rank."

Edge cases to volunteer:

Name the all-identical case. It's the cheapest check that duplicates are preserved, and a TreeSet-based solution fails it while passing almost everything else.

7. Follow-Up Questions — Modified Constraints

⭐ "Support the k-th SMALLEST instead."

Flip the heap: a max-heap of size k, whose root is the largest of the k smallest — again both the answer and the eviction candidate. The symmetry is exact, and being able to state it confirms you understood the inversion rather than memorised it.

⭐ "Allow k to change after construction."

Growing k is impossible — you already discarded the elements you'd now need. That's the cost of the space saving, and it's worth stating plainly. Shrinking k is easy: poll until the size matches. To support growth you'd have to retain everything, which puts you back to an order-statistic tree at O(log n) for arbitrary ranks.

"Support removing a previously added value."

PriorityQueue.remove(Object) is O(n) because a heap has no way to locate an arbitrary element. Options: a TreeMap multiset, which removes in O(log k); or lazy deletion — keep a count of pending removals and discard stale entries as they surface at the root. Lazy deletion is the standard trick and it's what question 6 uses in spirit.

"Return the top k values, not just the k-th."

The heap already holds exactly them — drain it, O(k log k), or iterate it unordered in O(k) if order doesn't matter. Note that iterating a PriorityQueue does not yield sorted order; only repeated poll does.

"What if k were close to n — say k = n/2?"

The heap holds half the stream, so space is O(n) and the log k is log n. At that point the space saving evaporates and sorting once is competitive. The heap wins when k ≪ n, which is worth saying rather than presenting it as universally better.

"What if the stream were distributed across machines?"

Each node keeps its own size-k min-heap, then the coordinator merges the k locals — because the global top k must be in the union of the per-node top ks. That's O(machines × k) to merge and it's the standard distributed top-k pattern.