Topics › Heaps › Hard Problems

Kth largest element in a stream of running integers

Hard TUF ID 567 SDE

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

Design a class that is constructed with an integer k and an initial array nums, and then receives more integers through add(val). After each add, return the k-th largest value among everything seen so far (the initial array plus all added values). “K-th largest” is by sorted position, not by distinct value: in [5, 5, 3] the 2nd largest is 5.

If fewer than k values have been seen, there is no k-th largest; in that case return -1. That is the convention of the GeeksforGeeks version, which starts from an empty stream and prints -1 for the first k - 1 numbers. LeetCode’s version guarantees at least k values whenever add is called, so the -1 branch never fires there.

Constraints (LeetCode): 1 ≤ k ≤ 10^4, 0 ≤ nums.length ≤ 10^4, values in [-10^4, 10^4], at most 10^4 calls to add.

Example 1

k = 3, nums = [4, 5, 8, 2]
add(3)  -> 4        // [2, 3, 4, 5, 8]: third largest is 4
add(5)  -> 5        // [2, 3, 4, 5, 5, 8]
add(10) -> 5        // [.., 5, 5, 8, 10]
add(9)  -> 8        // [.., 8, 9, 10]
add(4)  -> 8        // a small value does not change the top three

Example 2

k = 4, nums = []
add 1, 2, 3, 4, 5, 6  ->  -1, -1, -1, 1, 2, 3

The first three answers are -1 because fewer than four numbers exist; after that the answer is the 4th largest of the prefix.

All designs implement this interface:

interface KthLargest {
    int add(int val);
}

Approach 1: Brute force, keep everything and sort

Intuition. Store every number. When asked, sort in descending order and read index k - 1. It is easy to see that this is correct, and it makes a useful reference to test the faster designs against. But each add redoes work proportional to the whole history, even though only one number changed.

Algorithm.

1. The constructor stores k and all of nums.
2. add(val): append val. If the size is below k, return -1. Otherwise sort descending and return element k - 1.

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

    KeepAllSort(int k, int[] nums) {
        this.k = k;
        for (int x : nums) all.add(x);
    }

    public int add(int val) {
        all.add(val);
        if (all.size() < k) return -1;
        all.sort(Collections.reverseOrder());               // largest first
        return all.get(k - 1);
    }
}

Complexity. Each add is O(n log n) where n is the number of values seen. (Java’s TimSort actually runs close to linear here because the list is already almost sorted, but that is a library detail, not something to rely on in an interview.) Space O(n).

Approach 2: Better, keep only the top k in sorted order

Intuition. Here is the key observation behind every efficient solution: a value that drops out of the top k can never come back. The stream only ever adds numbers, so the k-th largest can only stay the same or increase. Any value already below it will stay below it forever. So we can throw such values away and keep at most k numbers.

Keep those k numbers in ascending order. The first element is then the smallest of the top k, which is exactly the k-th largest overall. A new value that is not bigger than that first element cannot enter the top k and is ignored. Otherwise we insert it at its sorted position and drop the first element if we now hold k + 1 values.

Algorithm.

1. insert(val): if we already hold k values and val ≤ top[0], return.
2. Binary search for the first index whose value is greater than val and insert there.
3. If the size exceeds k, remove index 0.
4. add(val): call insert, then return top[0] if we hold k values, else -1.

static class SortedTopK implements KthLargest {
    private final int k;
    private final List<Integer> top = new ArrayList<>();    // ascending, at most k values

    SortedTopK(int k, int[] nums) {
        this.k = k;
        for (int x : nums) insert(x);
    }

    private void insert(int val) {
        if (top.size() == k && val <= top.get(0)) return;   // cannot enter the top k
        int lo = 0, hi = top.size();                        // first index holding a value > val
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (top.get(mid) <= val) lo = mid + 1; else hi = mid;
        }
        top.add(lo, val);
        if (top.size() > k) top.remove(0);                  // drop the old k-th largest
    }

    public int add(int val) {
        insert(val);
        return top.size() < k ? -1 : top.get(0);
    }
}

Complexity. Binary search is O(log k), but inserting into or removing from the front of an ArrayList shifts up to k elements, so each add is O(k). Space O(k), independent of the length of the stream.

Approach 3: Optimal, a min-heap of size k

Intuition. Approach 2 keeps the top k fully sorted, yet only ever looks at the smallest of them. That is the job a min-heap does best: it keeps its minimum at the root and supports insert and remove-min in O(log k), without maintaining order among the other elements.

It may feel backwards to use a min-heap to find a largest element. The reason is that the heap holds the k largest values, and the boundary we care about, both for answering and for deciding what to evict, is the smallest of those k.

Algorithm.

1. offer(val): if the heap holds fewer than k values, push val.
2. Otherwise, if val > heap.peek(), pop the root (the old k-th largest) and push val. If val ≤ heap.peek(), ignore it.
3. The constructor offers every element of nums.
4. add(val): offer val; return heap.peek() if the heap holds k values, else -1.

static class MinHeapOfK implements KthLargest {
    private final int k;
    private final PriorityQueue<Integer> heap = new PriorityQueue<>();   // min-heap

    MinHeapOfK(int k, int[] nums) {
        this.k = k;
        for (int x : nums) offer(x);
    }

    private void offer(int val) {
        if (heap.size() < k) heap.offer(val);
        else if (val > heap.peek()) {                       // val beats the current k-th largest
            heap.poll();
            heap.offer(val);
        }
    }

    public int add(int val) {
        offer(val);
        return heap.size() < k ? -1 : heap.peek();
    }
}

Dry run on Example 1 (k = 3). The constructor offers 4, 5, 8 (heap fills up) and then 2, which is not above the root 4 and is ignored.

Complexity. Each add is O(log k): at most one pop and one push on a heap of size k. The constructor is O(m log k) for an initial array of length m. Space O(k).

Edge cases and pitfalls

• Fewer than k values. Calling peek() on an under-filled heap returns the k’-th largest for some smaller k', which is wrong. Check the size first and return -1 (or whatever the problem specifies).
• The -1 sentinel. If -1 can be a real answer, a sentinel is ambiguous. The tests avoid that clash, but in production code return an OptionalInt or throw instead.
• Duplicates count separately. [5, 5] with k = 2 has 2nd largest 5. Do not use a set.
• Strict versus non-strict comparison. Replacing the root only when val > root is correct and skips useless work when val == root, because the answer would not change.
• k = 1. The structure degenerates to a running maximum; the tests cover it with negative numbers.
• Initial array longer than k. The constructor must trim to k values, not blindly push everything, or later answers will be the wrong rank.

Related problems

• Kth largest element in a stream of running integers is the same problem as it appears in the heaps section.
• K-th largest element in an unsorted array. is the static version, where quickselect beats the heap.
• Find Median in a Stream tracks a moving boundary with two heaps instead of one.
• Top K Frequent Elements uses the same “min-heap of size k” trick on frequencies.
• Kth smallest element in an array [use priority queue] mirrors this with a max-heap.