[C++] 백준 2075번 N번째 큰 수

https://www.acmicpc.net/problem/2075

풀이

우선순위 큐를 사용하여 최소 힙에 입력을 받고, 힙의 크기가 n을 초과하면 pop하여, 힙에는 첫번째 큰 수부터 n번째 큰 수만 남도록 한다.

최종 n개가 힙에 남아있을 때, top이 n번째로 큰 수이다.

코드

#include <iostream>
#include <queue>
using namespace std;

int main()
{
    ios_base::sync_with_stdio(false);
	cin.tie(NULL); 
    cout.tie(NULL);
    
    int n;
    priority_queue<int, vector<int>, greater<int>> pq;
    
    cin >> n;
    
    for(int i=0; i<n*n; i++){
        int x;
        cin >> x;
        pq.push(x);
        
        if(pq.size() > n) pq.pop();
    }
    
    cout << pq.top();
    
    return 0;
}

입력 최적화를 하지 않으면 시간 초과가 발생한다!

---

다른 방법

우선순위 큐가 아닌 정렬을 사용하는 방법이다.

void nth_element(RandomIt first, RandomIt nth, RandomIt last);

nth_element(): n번째 원소를 찾고, n번째 원소보다 작은 값들은 앞쪽에, n번째 원소보다 큰 값들은 뒤쪽에 배치되므로 sort()보다 빠름

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    
    int n;
    cin >> n;
    
    vector<int> arr(n * n);
    for (int i = 0; i < n * n; i++) cin >> arr[i];

    nth_element(arr.begin(), arr.begin() + (n * n - n), arr.end());
    cout << arr[n * n - n];
    
    return 0;
}

소요 시간 및 메모리 비교

아래가 우선순위 큐를 이용한 방식, 위가 nth_element를 활용한 방식이다.

우선순위 큐는 시간이 조금 더 걸리지만 메모리를 훨씬 적게 사용한다.
nth_element 방식은 12MB의 메모리 제한에 가깝게 사용한다.