중앙값 구하기

백준 2696번: 중앙값 구하기

아이디어

최대 힙, 최소 힙 하나씩 준비한다.
1. 최대 힙에 원소 하나를 push한다.
2. 그 다음 원소부터는 최대 힙의 top에 위치한 원소보다 큰 경우는 최소 힙에 push하고, 아니면 최대 힙에 push한다.
3. 두개의 원소가 삽입된 후 최소 힙의 원소의 개수가 최대 힙의 원소의 개수보다 1개 많아질 때까지 각 힙에서 pop, push한다.
4. 최소 힙의 top에 있는 원소를 출력한다.

코드

#include <bits/stdc++.h>

using namespace std;

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int T;
    cin >> T;
    while (T--) {
        int M, x;
        priority_queue<int> pqM; //max
        priority_queue<int, vector<int>, greater<int>> pqm; //min
        
        cin >> M;
        cout << M/2+1 << '\n';
        cin >> x;
        pqM.push(x);
        cout << x << ' ';
        
        for (int i = 1; i < M; i++) {
            cin >> x;
            if (x > pqM.top()) pqm.push(x);
            else pqM.push(x);
            
            if (i%2 == 0) {
                if (pqM.size() >= pqm.size()) {
                    while (pqm.size()-pqM.size() != 1) {
                        pqm.push(pqM.top());
                        pqM.pop();
                    }
                }
                else {
                    while (pqm.size()-pqM.size() != 1) {
                        pqM.push(pqm.top());
                        pqm.pop();
                    }
                }
                if (i%20 == 0) cout << '\n';
                cout << pqm.top() << ' ';
            }
        }
        cout << '\n';
    }
    return 0;
}

여담

종만북에서 본듯한 문제다. 우선순위 큐 연습하기 좋은 문제.