[LeetCode] 703. Kth Largest Element in a Stream

https://leetcode.com/problems/kth-largest-element-in-a-stream/description/

Kth Largest Element in a Stream - LeetCode

문제 유형 : 우선순위 큐 / 힙

문제 난이도 : Easy

 

문제

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.

 

풀이

최소 힙 기반의 우선순위 큐를 생성하여 요소가 k개 미만이면 큐에 넣는다.

K개가 모이면, 그 뒤부터는 최소 힙의 맨 위보다 더 크면, 최소 힙의 맨위를 빼고,  새로 요소를 넣는다.

 

코드(C++)

class KthLargest {
public:
    int k = 0;
    priority_queue<int, vector<int>, greater<int>> pq;

    KthLargest(int k, vector<int>& nums) {
        this -> k = k;
        for(int i = 0; i < nums.size(); i++){
            if(i < k){
                pq.push(nums[i]);
            }
            else if(pq.top() < nums[i]){
                pq.pop();
                pq.push(nums[i]);
            }
        } 
    }
    
    int add(int val) {
        if(pq.size() < k){
            pq.push(val);
        }
        else if(pq.top() < val){
            pq.pop();
            pq.push(val);
        }
        return pq.top();
    }
};

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest* obj = new KthLargest(k, nums);
 * int param_1 = obj->add(val);
 */