[백준 2696번] 중앙값 구하기

문제 링크: https://www.acmicpc.net/problem/2696

2696번: 중앙값 구하기

<문제 풀이> 우선순위 큐

수를 입력받을 때마다 중앙값을 찾는 방법은 매 순간 정렬해서 가운데 수를 구하는 방법도 있지만 이는 매 순간 정렬하는데 O(nlogN)이 걸립니다. 최대 힙과 최소 힙을 이용하면 매 순간 정렬하지 않고 O(logN)에 중앙값을 구할 수 있습니다.

 

아이디어는 중앙값을 기준으로 왼쪽엔 중앙값보다 크기가 작거나 같은 수로 이루어진 최대 힙을 두고, 오른쪽엔 중앙값 보다 크기가 크거나 같은 수로 이루어진 최소 힙을 둡니다. 그리고 홀수 번째 수를 읽을 때마다 중앙값을 힙 사이즈가 작은 쪽에 push 한 뒤, 다시 힙 사이즈가 큰 것에 top을 중앙값으로 하면 됩니다. (중앙값을 기준으로 양 옆에 수의 개수를 일치시키는 느낌)

 

[알고리즘]

1. 최대 힙과 최소 힙을 준비한다.

2. 첫 번째 수를 중앙값으로 한다.

3. 두 번째부터 입력받은 수를 중앙값과 비교해서 작으면 최대 힙, 크면 최소 힙에 push한다.

4. 홀수 번째라면 최대 힙과 최소 힙 중에 사이즈가 작은 힙에 중앙값을 push 하고 그리고 힙 중에 사이즈가 큰 힙의 top을 중앙값으로 한다.

 

<C++ 소스 코드>

#include <iostream>
#include <queue>
using namespace std;
 
int main(void) {
    ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
 
    int test_case;
    cin >> test_case;
    while (test_case--) {
        int m;
        cin >> m;
        int data;
        int mid;
        priority_queue<int> max_pq;
        priority_queue<int, vector<int>, greater<int> > min_pq;
        vector<int> ans;
        cin >> data;
        mid = data;
        ans.push_back(mid); //1개 일때는 그 수가 mid
        for (int i = 2; i <= m; i++) {
            cin >> data;
            if (data < mid)max_pq.push(data); //data를 mid 보다 작으면 최대힙 크면 최소힙에 넣는다
            else min_pq.push(data);
            if (i % 2 == 1) { //홀수 번째 수를 읽을 때 마다
                int max_pq_size = max_pq.size();
                int min_pq_size = min_pq.size();
                if (max_pq_size < min_pq_size) { //최대 힙과 최소 힙 중에 사이즈가 작은 힙에 mid를 push
                    max_pq.push(mid);
                    max_pq_size++;
                }
                else {
                    min_pq.push(mid);
                    min_pq_size++;
                }
 
                if (max_pq_size > min_pq_size) { // 힙 중에 사이즈가 큰 힙의 top이 mid가 됨
                    mid = max_pq.top();
                    max_pq.pop();
                }
                else {
                    mid = min_pq.top();
                    min_pq.pop();
                }
                ans.push_back(mid);
            }
        }
        int size = ans.size();
        cout << size << "\n";
        for (int i = 1; i <= size; i++) {
            if (i != 1 && i % 10 == 1)cout << '\n';
            cout << ans[i-1] << " ";
        }
        cout << "\n";
    }
 
    return 0;
}