[C++/LeetCode] Kth Largest Element in a Stream

Daily LeetCoding Challenge for April, Day 8.

 

문제

배열에서 k번째로 큰 수를 찾음. 이때 add() 함수를 구현하여, 이 함수의 인자로 int 값을 주었을 때 배열에 추가하고 추가한 배열에서 k번째로 큰 수를 반환하도록 구현

 

아이디어: Heap(priority queue)

k번째로 큰 수를 찾는 것이므로, k+1번부터 큰 수는 필요하지 않다. 즉 배열에서 없어도 된다.또한 k번째로 큰 수를 찾는 것이므로 오름차순으로 정렬한 배열이 필요하다.이 조건을 쉽게 만족하는 자료구조는 바로 Heap이다.

 

Heap

1. complete binary tree2. 모든 노드의 값은 자식보다 크거나 같다. (max heap), 작거나 같다.(min heap)

 

Heap은 위의 특성을 만족하기 때문에 가장 큰 값, 가장 작은 값을 쉽게 구할 수 있다. 트리의 root가 바로 최대 또는 최소 값이기 때문이다.따라서 값을 넣으면 알아서 정렬해준다.c++에서는 priority queue가 heap을 써서 구현하였다. 

 

완성

submission: 32ms / 19.8MB

 

priority queue는 kth 큰 수이기 때문에 min heap으로 구현한다.

먼저 생성자에서는 배열의 값을 저장하는데, 이때 pq.size()가 k보다 큰 경우, 즉 heap에 k보다 많은 수가 들어온 경우, 저장할 필요가 없다. 따라서 이는 tree에 넣어서 정렬을 한 뒤 pop() 을 통해 제거한다.

왜냐하면 k+1개의 node가 heap에 있는 경우, tree의 top()은 k+1 번째로 큰 수이기 때문이다. (min heap이라서)

 

그 후 add 함수에도 똑같은 논리를 적용해주면 된다.

 

class KthLargest {
private:
    priority_queue<int, vector<int>, greater<int>> pq;
    int index;
public:
    KthLargest(int k, vector<int>& nums) {
        for(int i = 0;i < nums.size();i++)
        {
            pq.push(nums[i]);
            if (pq.size() > k)
                pq.pop();
        }
        index = k;
    }
    
    int add(int val) {
        pq.push(val);
        if (pq.size() > index)
            pq.pop();
        return pq.top();
    }
    
};