Topics › Heaps › Hard Problems

Kth largest element in a stream of running integers

Hard TUF ID 2838 A2Z

Problem

Design a class that is created with an integer k and an initial array nums. It then receives a stream of integers through add(val). After each add, return the k-th largest value among everything seen so far, the initial array included. Duplicates count separately, so in [7, 7, 8] the 2nd largest is 7.

In the LeetCode version, at least k values are present whenever add returns. The GeeksforGeeks version allows fewer and expects -1 in that case, so the code here returns -1 while fewer than k values exist and handles both versions.

Constraints: 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      // values 8 5 4 3 2: the 3rd largest is 4
add(5)  -> 5      // 8 5 5 4 3 2
add(10) -> 5      // 10 8 5 5 ...
add(9)  -> 8      // 10 9 8 5 ...
add(4)  -> 8

Example 2

k = 3, nums = []
add(5) -> -1      // only one value so far
add(1) -> -1
add(2) -> 1       // values 5 2 1: the 3rd largest is 1
add(7) -> 2

All three designs implement one interface, so one test script drives them all:

interface KthStream {
    int add(int val);
}

Approach 1: Brute force, sort on every add

Intuition. Keep every value in a list. When asked, sort the values and read the k-th from the end. This follows the definition directly and is clearly correct. It is also slow, because it re-sorts everything for every single add.

Algorithm.

  1. The constructor stores k and every value of nums.
  2. add(val): append val. If fewer than k values exist, return -1. Otherwise copy the values into an array, sort it, and return a[size - k].
static class SortEveryTime implements KthStream {
    private final int k;
    private final List<Integer> all = new ArrayList<>();

    SortEveryTime(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;
        int[] a = new int[all.size()];
        for (int i = 0; i < a.length; i++) a[i] = all.get(i);
        Arrays.sort(a);
        return a[a.length - k];
    }
}

Complexity. Time O(n log n) per add, where n is the number of values so far. Space O(n).

With 2 · 10^4 values and 10^4 adds, that is billions of comparisons. Notice how much of the sorted array is wasted: we read exactly one position.

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

Intuition. The key observation for every faster approach: a value that is not among the k largest now can never be the k-th largest later. The stream only grows. A new value can push existing values down in rank, never up. So once a value drops out of the top k, we can forget it for good.

Keep just the k largest values in an ascending list. Its first element is the smallest of the top k, which is the k-th largest overall. For a new value, find its position with binary search, insert it, and if the list now holds k + 1 values, remove the first one.

Algorithm.

  1. The constructor stores k and calls add for each initial value, ignoring the returned answers.
  2. add(val): binary search for the first index whose value is greater than val and insert there. If the size exceeds k, remove index 0.
  3. Return -1 if the size is still below k, else top.get(0).
static class SortedTopK implements KthStream {
    private final int k;
    private final List<Integer> top = new ArrayList<>();      // ascending; top.get(0) is the k-th largest once full

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

    public int add(int val) {
        int lo = 0, hi = top.size();                          // find the first index whose value is > 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 smallest: it can never be k-th largest again
        return top.size() < k ? -1 : top.get(0);
    }
}

Complexity. Time O(k) per add. The binary search is O(log k), but inserting into and removing from the front of an ArrayList shifts up to k elements. Space O(k).

This is a real improvement: the cost no longer grows with the stream length. The remaining waste is that we keep the entire top k in sorted order when we only ever read its minimum.

Approach 3: Optimal, min-heap of size k

Intuition. We need three things from our collection of the k largest values: its minimum (the answer), the ability to remove that minimum when a bigger value arrives, and the ability to insert. A min-heap does all three, with peek in O(1) and poll and add in O(log k). Full sorted order is unnecessary.

It can feel backwards to use a min-heap for a “largest” question. The heap does not hold all values, only the top k, and among those we care about the smallest, because that one is the boundary. A new value either beats the boundary and enters, pushing the old boundary out, or it does not and changes nothing.

Algorithm.

  1. The constructor stores k and calls add for each initial value.
  2. add(val):
    • If the heap has fewer than k values, push val.
    • Else, if val > heap.peek(), poll the minimum and push val. If val is not larger, it cannot be in the top k, so ignore it.
  3. Return -1 if the heap has fewer than k values, else heap.peek().
static class MinHeapK implements KthStream {
    private final int k;
    private final PriorityQueue<Integer> heap = new PriorityQueue<>();

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

    public int add(int val) {
        if (heap.size() < k) {
            heap.add(val);
        } else if (val > heap.peek()) {                       // val beats the current k-th largest
            heap.poll();
            heap.add(val);
        }
        return heap.size() < k ? -1 : heap.peek();
    }
}

Dry run on example 1, k = 3:

call heap after (the top k) returned
constructor: 4, 5, 8 {4, 5, 8}  
constructor: 2 2 < 4, ignored: {4, 5, 8}  
add(3) 3 < 4, ignored: {4, 5, 8} 4
add(5) 5 > 4: drop 4, {5, 5, 8} 5
add(10) 10 > 5: drop 5, {5, 8, 10} 5
add(9) 9 > 5: drop 5, {8, 9, 10} 8
add(4) 4 < 8, ignored 8

Complexity. Time O(log k) per add, and O(n log k) for the constructor. Space O(k), independent of how long the stream runs.

Edge cases and pitfalls