703. Kth Largest Element in a Stream

1. Description

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.
Implement KthLargest class:

• KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums.
• int add(int val) Appends the integer val to the stream and returns the element representing the kth largest element in the stream.

2. Example

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

3. Constraints

• 1 <= k <= $10^4$
• 0 <= nums.length <= $10^4$
• $-10^4$ <= nums[i] <= $10^4$
• $-10^4$ <= val <= $10^4$
• At most $10^4$ 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 kth element.

4. Solutions

Heap(Priority Queue)

class KthLargest {
public:
    KthLargest(int k, vector<int> &nums) {
    // n = nums.size()  
    // Time complexity: O(nlogk)  
    // Space complexity: O(k)
        k_ = k;
        for (int i = 0; i < nums.size(); ++i) {
            if (data_.size() < k) {
                data_.push(nums[i]);
            } else if (nums[i] > data_.top()) {
                data_.pop();
                data_.push(nums[i]);
            }
        }
    }

    int add(int val) {
    // Time complexity: O(logk)  
    // Space complexity: O(1)
        if (data_.size() < k_) {
            data_.push(val);
            return data_.top();
        } else if (val < data_.top()) {
            return data_.top();
        } else {
            data_.pop();
            data_.push(val);
            return data_.top();
        }
    }

private:
    int k_;
    priority_queue<int, vector<int>, greater<int>> data_;
};