[BOJ] 2696번 - 중앙값 구하기

백준 2696번 - 중앙값 구하기

시간제한	메모리 제한	제출	정답	맞은 사람	정답 비율
1 초	128MB	3196	1608	1280	52.697%

문제

어떤 수열을 읽고, 홀수번째 수를 읽을 때 마다, 지금까지 입력받은 값의 중앙값을 출력하는 프로그램을 작성하시오.

예를 들어, 수열이 1, 5, 4, 3, 2 이면, 홀수번째 수는 1번째 수, 3번째 수, 5번째 수이고, 1번째 수를 읽었을 때 중앙값은 1, 3번째 수를 읽었을 때는 4, 5번째 수를 읽었을 때는 3이다.

입력

첫째 줄에 테스트 케이스의 개수 T(1 ≤ T ≤ 1,000)가 주어진다. 각 테스트 케이스의 첫째 줄에는 수열의 크기 M(1 ≤ M ≤ 9999, M은 홀수)이 주어지고, 그 다음 줄부터 이 수열의 원소가 차례대로 주어진다. 원소는 한 줄에 10개씩 나누어져있고, 32비트 부호있는 정수이다.

 

출력

각 테스트 케이스에 대해 첫째 줄에 출력하는 중앙값의 개수를 출력하고, 둘째 줄에는 홀수 번째 수를 읽을 때 마다 구한 중앙값을 차례대로 공백으로 구분하여 출력한다. 이때, 한 줄에 10개씩 출력해야 한다.

---

예제 입력 1

3
9
1 2 3 4 5 6 7 8 9
9
9 8 7 6 5 4 3 2 1
23
23 41 13 22 -3 24 -31 -11 -8 -7
3 5 103 211 -311 -45 -67 -73 -81 -99
-33 24 56

예제 출력 1

5
1 2 3 4 5
5
9 8 7 6 5
12
23 23 22 22 13 3 5 5 3 -3
-7 -3

출처

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

2696번: 중앙값 구하기

알고리즘 분류

• 자료구조
• 우선순위 큐

접근 방법

이 문제를 푸는 가장 Naive 한 방법은 원소가 추가 될 때 마다, 정렬을 수행하여 중앙값을 찾는 것이다.

비교기반의 정렬 함수 중 가장 빠른 것을 이용하여 풀어도, 시간복잡도는 O( N * NlogN ) 이고 TLE이다.

 

정렬에서 좀 더 최적화 할 수 있는 부분은 없는지, 혹은 정렬을 수행하지 않고도 중앙값을 찾는 방법이 없는지 통찰을 얻어보자. 

먼저, 특정 N번째 원소까지의 중앙 값을 ${A_k}$라고 가정한다면, 우리는 다음과 같이 이분화를 할 수 있게 된다.

${A_1 ... A_n}$ = { ${A_k}$ 보다 작은 원소들의 집합 } + { ${A_k}$ 보다 크거나 같은 원소들의 집합 }

이 때, n이 홀수 일 때 유일한(unique) 중앙값이 하나가 정의가 되고 이 때 이것을 출력하고자 한다.

이 경우에 { ${A_k}$ 보다 작은 원소들의 집합 } + 1 이  { ${A_k}$ 보다 크거나 같은 원소들의 집합 }의 개수일 것이다. 그리고 { ${A_k}$ 보다 크거나 같은 원소들의 집합 } 중 최솟값이 중앙값이 될 것이다.

 

위에서 설명한 내용의 핵심은 3가지이다.

1. { ${A_k}$ 보다 작은 원소들의 집합 } 과 { ${A_k}$ 보다 크거나 같은 원소들의 집합 } 은 최대 1이 차이나고, 항상 { ${A_k}$ 보다 크거나 같은 원소들의 집합 }은 { ${A_k}$ 보다 작은 원소들의 집합 } 보다 크거나 같다.
2. { ${A_k}$ 보다 작은 원소들의 집합 } 의 모든 원소들은  { ${A_k}$ 보다 크거나 같은 원소들의 집합 } 의 모든 원소들 보다 작다.
   1. 즉, { ${A_k}$ 보다 작은 원소들의 집합 }의 최댓값이 { ${A_k}$ 보다 크거나 같은 원소들의 집합 }의 최솟값 보다 작아야 한다는 것을 의미한다.
3. ${A_1 ... A_n}$ 을 위와 같이 이분화 하는 방법은 Unique 하다.
   1. 즉, 이것이 중앙값을 구할 수 있는 유일한 해이다. 

 

이러한 성질을 어떠한 자료구조로 구현하는 것이 적절한지 감이 올 것이다.

바로 Priority Queue 이다. 

{ ${A_k}$ 보다 작은 원소들의 집합 } (이하, lower) 은 최댓값을 빠르게 구하기 위해 Max heap으로 구현을 하고

{ ${A_k}$ 보다 크거나 같은 원소들의 집합 } (이하, upper)은 최솟값을 빠르게 구하기 위해 Min heap으로 구현을 하자.

 

원소가 추가될 때, 우리는 lower든, upper든 아무곳에나 삽입해도 상관이 없다. 코드에서는 upper에 삽입한 다음 적당히 adjusting하여 앞에서 설명한 2가지 특성들을 만족하도록 유지시킬 것이다.

그렇게 유지시킨다면, upper의 원소가 lower의 원소보다 1개 더 많을 때, upper의 최솟값이 현재까지의 중앙값이다.

 

특성 유지( Adjusting ) 

1. { ${A_k}$ 보다 작은 원소들의 집합 } 과 { ${A_k}$ 보다 크거나 같은 원소들의 집합 } 은 최대 1이 차이나고, 항상 { ${A_k}$ 보다 크거나 같은 원소들의 집합 }은 { ${A_k}$ 보다 작은 원소들의 집합 } 보다 크거나 같다.
2. { ${A_k}$ 보다 작은 원소들의 집합 } 의 모든 원소들은  { ${A_k}$ 보다 크거나 같은 원소들의 집합 } 의 모든 원소들 보다 작다.
   1. 즉, { ${A_k}$ 보다 작은 원소들의 집합 }의 최댓값이 { ${A_k}$ 보다 크거나 같은 원소들의 집합 }의 최솟값 보다 작아야 한다는 것을 의미한다.

위의 특성을 유지하기 위해 아래 2,3번 같은 과정을 거치면 된다.

 

1. upper에 새로 들어 온 원소 삽입 (이때, 특성이 깨질 수 있다)
2. upper의 최솟값을 lower로 옮긴다.(최소한, 특성2을 유지)
3. 특성1에 충족하지 않을 경우(lower.size() > upper.size()), lower의 최댓값을 upper로 옮긴다.(특성 1을 유지)

 

소스 코드

#include <bits/stdc++.h>
using namespace std;
int main(void){
    ios_base::sync_with_stdio(false);
    cin.tie(0), cout.tie(0);

    int T;
    cin >> T;
    while(T--){
        int M;
        cin >> M;
        priority_queue<int, vector<int>> lower; // -> maxHeap
        priority_queue<int, vector<int>, greater<>> upper; // -> minHeap

        if( M & 1 ) cout << M/2 + 1 << '\n';
        else cout << M/2 << '\n';

        for(int i=0; i<M; i++){
            int elem;
            cin >> elem ;
            upper.push(elem);
            lower.push(upper.top()); upper.pop();
            if( lower.size() > upper.size() ) {
                upper.push(lower.top()); lower.pop();
            }
            if( i % 2 == 0 )
                cout << upper.top() << ' ';
        }
        cout << '\n';
    }

    return 0;
}