Kth Largest Element in a Stream

Description

Kth Largest Element in a Stream

You are part of a university admissions office and need to keep track of the kth highest test score from applicants in real-time. This helps to determine cut-off marks for interviews and admissions dynamically as new applicants submit their scores.

You are tasked to implement a class which, for a given integer k, maintains a stream of test scores and continuously returns the kth highest test score after a new score has been submitted. More specifically, we are looking for the kth highest score in the sorted list of all scores.

Implement the KthLargest class:

• KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of test scores nums.
• int add(int val) Adds a new test score val to the stream and returns the element representing the kth largest element in the pool of test scores so far.

Example 1:
Input:
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output: [null, 4, 5, 5, 8, 8]
Explanation:
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3); // return 4
kthLargest.add(5); // return 5
kthLargest.add(10); // return 5
kthLargest.add(9); // return 8
kthLargest.add(4); // return 8

Example 2:
Input:
["KthLargest", "add", "add", "add", "add"]
[[4, [7, 7, 7, 7, 8, 3]], [2], [10], [9], [9]]
Output: [null, 7, 7, 7, 8]

Explanation:
KthLargest kthLargest = new KthLargest(4, [7, 7, 7, 7, 8, 3]);
kthLargest.add(2); // return 7
kthLargest.add(10); // return 7
kthLargest.add(9); // return 7
kthLargest.add(9); // return 8

Constraints:

• 0 <= nums.length <= 104
• 1 <= k <= nums.length + 1
• -104 <= nums[i] <= 104
• -104 <= val <= 104
• At most 104 calls will be made to add.

Primary Approach: Min-Heap of Size $k$ ($O(\log k)$ Time per Add, $O(k)$ Space)

Intuition

To find the $k$-th largest element in a stream, we only care about maintaining the top $k$ largest numbers seen so far. A Min-Heap (PriorityQueue) of size $k$ is the ideal data structure:

1. Min-Heap Property: The smallest element among the top $k$ elements always stays at the root (minHeap.peek()).
2. $k$-th Largest Definition: The smallest of the $k$ largest elements is precisely the $k$-th largest element overall.
3. Execution:
   • add(val): Push val into minHeap. If minHeap.size() exceeds $k$, evict the smallest element (minHeap.poll()).
   • Return minHeap.peek(), which guarantees the $k$-th largest value in $O(1)$ lookup time.

import java.util.PriorityQueue;
 
class KthLargest {
    private final PriorityQueue<Integer> minHeap;
    private final int k;
 
    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.minHeap = new PriorityQueue<>();
 
        // Add initial array elements to the stream
        for (int num : nums) {
            add(num);
        }
    }
    
    public int add(int val) {
        minHeap.offer(val);
 
        // Maintain heap size of at most k
        if (minHeap.size() > k) {
            minHeap.poll();
        }
 
        return minHeap.peek();
    }
}
 

Complexity

• Time Complexity:
  • Constructor: $O(N\log k)$ — Inserting $N$ initial elements into a heap bounded by size $k$.
  • add(val): $O(\log k)$ — Pushing and popping from a heap containing at most $k+1$ elements.
• Space Complexity: $O(k)$ — Min-heap stores at most $k$ elements at any point.

---

Key Interview Discussion Points

• Why Min-Heap instead of Max-Heap?

• A Max-Heap would require storing all stream elements ($N$ elements), taking $O(N)$ space and $O(\log N)$ time per add() call.

• A Min-Heap of size $k$ discards all elements smaller than the $k$-th largest, reducing space to $O(k)$ and operation time to $O(\log k)$, which is significantly faster for large streams ($N\gg k$).

• Handling Initial Array Edge Case: The problem allows nums.length < k initially. By routing initial items through add(), minHeap.poll() will only trigger once the heap size exceeds $k$, ensuring safety when nums starts out empty or small.

---

Easy Memory Rule

“Finding $k$-th LARGEST? Use a MIN-HEAP of size $k\to$ minHeap.peek() always gives the answer!”