[백준/C++] #1696 중앙값 구하기 (자료구조)

문제

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

 

구상

정렬을 적용한 배열(벡터)이나 우선순위 큐를 사용하여 해당 상황에서의 중앙값 판별해야겠다 싶었다. 
그런데, 새로운 원소가 들어올 때마다 벡터를 정렬시키는 방법이랑 우선순위 큐 하나를 이용하여 push, pop을 반복하는 방법은 시간복잡도가 너무 컸다. 

 

우선순위 큐 하나를 사용하여 push, pop을 반복하며 중앙값을 구한다고 하면 최악의 경우 O(T * M * M logM)이 걸린다. 

*우선순위 큐에서 요소 하나 삽입하거나 삭제할 때 O(logN) 소요.

*요소 M개라 가정하면 테케 하나에 전체 수열의 길이가 M이라면, 중앙값들 구할때 O(M *  M logM) 소요.

(원소 하나 추가할 때마다 pop 반복하며 중앙값 구해야 함 * 원소 개수)

*테스트 케이스 수만큼 반복하면 O(T * M * M logM) 소요.

=> 문제에 적용을 하면 1000*9999*9999*log9999 => 1억 초과 => 시간초과

그렇다면 어떤방법을 사용해야할까?

힙 하나로는 안되니까, 힙 두개를 사용해볼까? max heap에는 중앙값보다 작은 수들을 넣고, min heap에는 중앙값보다 큰 수들을 담아볼까? 두개 힙의 크기만 잘 조정을 시키면 힙의 루트값이 중앙값이 되니까 구하기 쉽겠다고 판단했다. 이 경우에는, 전체 수열의 길이가 M이라면 O(T * M logM)가 소요되어 1억번 이하의 연산이 소요된다. 

 

풀이 + 코드

#include <iostream>
#include <queue>
#include <vector>
using namespace std;
/*
2696 중앙값 구하기
어떤 수열을 읽고, 홀수번째 수를 읽을 때마다, 지금까지 입력받은 값의 중앙값을 출력
*/

void findMedians(int M, vector<int>& sequence) {
    priority_queue<int> maxHeap; // 최대 힙 (left side)
    priority_queue<int, vector<int>, greater<int>> minHeap; // 최소 힙 (right side)

    vector<int> medians; // 중앙값 저장

    for (int i = 0; i < M; ++i) {
        // 힙에 요소 넣기
        if (maxHeap.empty() || sequence[i] <= maxHeap.top()) {
            maxHeap.push(sequence[i]);
        }
        else {
            minHeap.push(sequence[i]);
        }

        // 힙의 균형 유지 (maxHeap이 하나 더 많거나 같도록)
        if (maxHeap.size() > minHeap.size() + 1) { // 현재: maxHeap.size() == minHeap.size() + 2
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        }
        else if (minHeap.size() > maxHeap.size()) { // 현재: minHeap.size() == maxHeap.size() + 1
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }

        // 홀수 번째 수일 때 중앙값 저장
        if (i % 2 == 0) {
            medians.push_back(maxHeap.top());
        }
    }

    // 출력
    cout << medians.size() << endl;
    for (int i = 0; i < medians.size(); ++i) {
        if (i > 0 && i % 10 == 0) cout << endl;
        cout << medians[i] << " ";
    }
    cout << endl;
}

int main() {
    int T;
    cin >> T;

    while (T--) {
        int M;
        cin >> M;
        vector<int> sequence(M);

        for (int i = 0; i < M; ++i) {
            cin >> sequence[i];
        }

        findMedians(M, sequence);
    }

    return 0;
}