[백준,c++] 2696번 - 중앙값 구하기

문제

2696번: 중앙값 구하기

 

코드

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

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    
    int T; cin>>T;
    while(T--){
        int n; cin>>n;
        priority_queue<int> max_heap;   // max_heap에는 중간값 이하의 값만 담겨야 한다.
        priority_queue<int,vector<int>,greater<>> min_heap;
        vector<int>ans;
        int mid = 0;
        for(int i=1; i<=n; i++){
            int number; cin>>number;
            if(i==1) mid = number;
            
            if(mid<number) min_heap.push(number);
            else max_heap.push(number);
            
            if(i&1){    // 홀수 번째 숫자일 때
                while(min_heap.size()>max_heap.size()){
                    mid = min_heap.top();
                    max_heap.push(mid);
                    min_heap.pop();
                }
                while(min_heap.size()<max_heap.size()){
                    mid = max_heap.top();
                    min_heap.push(mid);
                    max_heap.pop();
                }
                ans.push_back(mid);
            }
            
        }
        
        cout<<ans.size()<<'\n';
        for(int i=0; i<ans.size(); i++){
            if(i%10==0&&i!=0) cout<<'\n';
            cout<<ans[i]<<' ';
        }
    }
}

 

풀이

상당히 어려웠고 풀이와 이해까지 오래 걸렸습니다. 일단 이 문제의 경우 테스트 케이스가 1000, 수열의 크기가 9999까지 들어옵니다. 따라서 매번 입력값이 들어올 때마다 정렬한 후 중앙값을 뽑아내게 되면 1000*9999*정렬(nlogn)이기 때문에 제한시간 안에 문제를 해결하지 못합니다. 따라서 max_heap과 min_heap 이중 우선순위 큐를 써서 문제를 구현해야 합니다.

문제 핵심 해결방법은 아래 세 가지입니다.

1. max_heap에는 중앙값 이하의 값만 들어갈 수 있다.
2. 홀수번째 숫자일 때 max_heap과 min_heap의 개수를 동일하게 맞춰야 한다.(많은 쪽에서 적은 쪽으로 이동한다.)
3. 2번에서 이동할 때마다 이동하는 값이 중앙값이 되기 때문에 계속 갱신해 준다.

다음번에 유사한 문제가 있을 때는 이번 문제를 떠올려서 손쉽게 풀 수 있었으면 좋겠습니다.