[백준] 1655번 : 가운데를 말해요

문제

 

---

풀이

 

이 문제의 시간 제한은 0.1초다.

그니까 입력이 들어오자마자 바로바로 출력을 내보내야 한다는 것이다.

입력을 받을 때를 제외하곤 단 하나의 for문도 쓸 수 없다.

 

그럼 도대체 제일 작은 값도 제일 큰 값도 아닌 중앙값을 바로바로 내보낼 것인가...

중앙값을 기준으로 그보다 작은 수는 최대 힙에 넣고 그보다 큰 수는 최소 힙에 넣는 것이다.

그렇다면 최대 힙의 top에는 중앙값보다 같거나 작은 수 중 가장 큰 수가,

최소 힙의 top에는 중앙값보다 같거나 큰 수 중 가장 작은 수가 올 것이다.

 

균형을 이루기 위해선 최대 힙과 최소 합의 원소의 갯수 차이가 0(원소의 총 갯수가 홀수)이거나 1(원소의 총 갯수가 짝수)이여야 한다.

만약 이 균형이 깨진다면 중앙값을 수정해야 한다.

 

위 상태에서 10과 11이 들어온 경우를 생각해보자

10이 들어올 때까지는 두 힙의 크기 차이가 1이기 때문에 5를 여전히 중앙값으로 쓸 수 있지만, 11이 들어오면 균형이 깨진다.

이때 최소 힙의 top에 있는 6을 새로운 중앙값으로 바꿔주고 기존의 중앙값이던 5는 최대 힙으로 넣어준다.

 

 

이번엔 0이 들어왔다고 하자.

얼핏봐선 균형이 깨지지 않는 것 같지만 그렇지 않다.

중앙값의 후보가 둘일 때는 둘 중 더 작은 값을 선택해야 한다. 이 규칙에 따르면 중앙값은 5가 된다.

그렇기 때문에 최소 힙의 크기가 최대 힙의 크기보다 큰 경우에는 그 차이가 2가 되어야 중앙값의 변화가 생기는 것과 달리,

그 반대의 경우에는 그 차이가 1만 되어도 중앙값에 변화가 생긴다.

 

따라서 최대 힙의 top에 있던 5를 새로운 중앙값으로 하고, 기존의 중앙값이던 6은 최소 힙에 넣는다.

---

소스코드

 

#include <iostream>
#include <queue>

using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);

    priority_queue<int> max_pq;
    priority_queue<int, vector<int>, greater<int>> min_pq; //greater면 min heap
    int N, x, mid;

    cin >> N;
    cin >> mid; //첫 입력은 중앙값
    cout << mid << '\n';
    for (int i = 1; i < N; i++) {
        cin >> x;
        if (x > mid) //중앙값보다 큰 값은 최소 힙으로
            min_pq.push(x);
        else //중앙값보다 작거나 같은 값은 최대 힙으로
            max_pq.push(x);

        int interval = max_pq.size() - min_pq.size(); //두 우선순위 큐의 크기 차이는 항상 0 또는 1이어야 함.
        if (interval > 0) { //최대 힙의 원소가 더 많은 경우
            min_pq.push(mid);
            mid = max_pq.top();
            max_pq.pop();
        } else if (interval < -1) { //최소 힙의 원소가 더 많은 경우
            max_pq.push(mid);
            mid = min_pq.top();
            min_pq.pop();
        }
        cout << mid << '\n';
    }
}

전체 소스코드다.

 

    cin >> N;
    cin >> mid; //첫 입력은 중앙값
    cout << mid << '\n';

첫번째 입력을 첫번째 중앙값으로 초기화 한다.

 

        cin >> x;
        if (x > mid) //중앙값보다 큰 값은 최소 힙으로
            min_pq.push(x);
        else //중앙값보다 작거나 같은 값은 최대 힙으로
            max_pq.push(x);

그리고 중앙값보다 큰 값은 최소 힙으로, 작거나 같은 값은 최대 힙으로 넣는다.

 

        int interval = max_pq.size() - min_pq.size(); //두 우선순위 큐의 크기 차이는 항상 0 또는 1이어야 함.
        if (interval > 0) { //최대 힙의 원소가 더 많은 경우
            min_pq.push(mid);
            mid = max_pq.top();
            max_pq.pop();
        } else if (interval < -1) { //최소 힙의 원소가 더 많은 경우
            max_pq.push(mid);
            mid = min_pq.top();
            min_pq.pop();
        }
        cout << mid << '\n';

힙에 입력을 넣은 뒤, 각 힙의 크기를 비교한다.

최대 힙의 원소가 최소 힙보다 1개라도 더 많다면 바로 중앙값을 변경하고,

최소 힙의 원소가 더 많은 경우에는 그 차이가 2가 될 때 중앙값을 변경한다.