Kth Largest Element In a Stream

Easy Top 250

Associated Patterns

Interviewed At (Company Tags)

GoogleAmazonMicrosoftMeta

Problem Description

Design a class to find the k-th largest element in a stream. Note that it is the k-th largest element in the sorted order, not the k-th distinct element.

Implement KthLargest class:

KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of numbers nums.
int add(int val) Appends the integer val to the stream and returns the element representing the k-th largest element.

Examples

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

Constraints

1 <= k <= 10⁴
0 <= nums.length <= 10⁴
-10⁴ <= nums[i] <= 10⁴
-10⁴ <= val <= 10⁴
At most 10⁴ calls will be made to add.
It is guaranteed that there will be at least k elements in the array when you search for the k-th element.

Min-Heap of Size K

To find the k-th largest element efficiently in a stream, we can maintain a Min-Heap of size k.

The min-heap stores the k largest elements seen so far.
The root of the min-heap represents the minimum of these k elements, which is the k-th largest element. When a new element val is added:

Push val into the heap.
If the heap size exceeds k, pop the smallest element from the top of the heap.
The element at the top of the heap is the result. All insertions execute in O(log k) time.

Solution 1: Min-Heap Tracker of Size K

Initialize a min-heap, maintaining its size at k elements on each insertion.

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<>();
        for (int num : nums) {
            add(num);
        }
    }

    public int add(int val) {
        minHeap.offer(val);
        if (minHeap.size() > k) {
            minHeap.poll(); // discard smallest element
        }
        return minHeap.peek(); // top of min-heap of size k is the kth largest
    }
}

import heapq

class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)  # discard smallest
        return self.heap[0]  # root is the kth largest

#include <vector>
#include <queue>

class KthLargest {
private:
    std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; // min-heap
    int k;

public:
    KthLargest(int k, std::vector<int>& nums) {
        this->k = k;
        for (int num : nums) {
            add(num);
        }
    }

    int add(int val) {
        minHeap.push(val);
        if (minHeap.size() > k) {
            minHeap.pop();
        }
        return minHeap.top();
    }
};

Complexity Analysis

Time Complexity: Constructor takes O(n log k) time. The add function takes O(log k) time.
Space Complexity: O(k) auxiliary space to store elements in the heap.