[Algospot] 변화하는 중간값 (문제 ID : RUNNINGMEDIAN)

문제

한 수열의 중간값(median)은 이 수열을 정렬했을 때 가운데 오는 값입니다. 예를 들어 {3,1,5,4,2}를 정렬했을 때 가운데 오는 값은 3이지요. 수열의 길이가 짝수일 때는 가운데 있는 두 값 중 보다 작은 것을 수열의 중간값이라고 정의하도록 합시다.

한 수열의 중간값은 수열에 새로운 수가 추가될 때마다 바뀔 수 있습니다. 텅 빈 수열에서 시작해서 각 수가 추가될 때마다 중간값을 계산하는 프로그램을 작성하세요. 예를 들어 3, 1, 5, 4, 2 순서대로 숫자가 추가될 경우 수열의 중간값은 3, 1, 3, 3, 3 순서로 변화합니다.

입력 생성

입력의 크기가 큰 관계로, 이 문제에서는 수열을 입력받는 대신 다음과 같은 식을 통해 프로그램 내에서 직접 생성합니다.

• A[0] = 1983

• A[i] = (A[i-1] * a + b) % 20090711

a와 b는 입력에 주어지는 상수입니다. 이 문제의 해법은 입력을 생성하는 방식과는 아무 상관이 없습니다.

 

입력

입력 파일의 첫 줄에는 테스트 케이스의 수 C (1 <= C <= 20)가 주어지고, 그 후 C줄에 각 3개의 정수로 수열의 길이 N (1 <= N <= 200,000), 그리고 수열을 생성하는 데 필요한 두 정수 a , b (0 <= a,b <= 10000) 가 주어집니다.

 

출력

각 테스트 케이스마다 한 줄에 N개의 중간값의 합을 20090711로 나눈 나머지를 출력합니다.

---

중간값을 구해 모두 더하는 문제이다. 정답 코드는 힙을 이용해서 푸는 방식인데, 문제를 읽고 시간초과가 날 것을 알고 내 방식으로 문제를 풀었다.

역시나 시간초과가 나왔다. 그냥 배열에 값을 하나씩 넣으면서 STL sort 함수를 써서 그때 그때 중간값을 더했는데, 역시 sort도 오래 걸릴 것이고 시간 초과가 떴다.

 

정답 방식은 숫자들을 정렬한 뒤 앞의 절반을 최대힙에, 뒤의 절반을 최소힙에 넣는것이다. 그렇다면 최대힙의 루트가 중간값이 될 것이다. 이를 두개의 불변식으로 표현한다.

1. 최대힙의 크기는 최소힙의 크기와 같거나 하나 더 크다.

2. 최대힙의 최대 원소는 최소힙의 최소 원소보다 작거나 같다.

이를 만족하는 방식으로 문제를 해결한다.

 

일단 1번식을 만족하는 방식으로, 두 힙의 크기에 따라 seed를 삽입한다.

그 후엔 2번식을 만족시키게 두 힙의 루트를 비교한다. 

 

코드 원본 : https://github.com/chosh95/STUDY/blob/master/Algospot/RUNNINGMEDIAN.cpp

chosh95/STUDY

C++ 코드

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std; 
int C, N, a, b;

int median()
{
	long long seed = 1983;
	priority_queue<int, vector<int>, less<int>> maxHeap;
	priority_queue<int, vector<int>, greater<int>> minHeap;
	int medSum = 0;
	for (int i = 1; i <= N; i++) {
		if (maxHeap.size() == minHeap.size())
			maxHeap.push(seed);
		else
			minHeap.push(seed);
		if (!minHeap.empty() && !maxHeap.empty() && minHeap.top() < maxHeap.top()) {
			int tmpA = maxHeap.top(), tmpB = minHeap.top();
			maxHeap.pop();
			minHeap.pop();
			maxHeap.push(tmpB);
			minHeap.push(tmpA);
		}
		medSum = (medSum + maxHeap.top()) % 20090711;
		seed = (seed * a + b) % 20090711;
	}
	return medSum;
}

int main()
{
	cin >> C;
	while (C--) {
		cin >> N >> a >> b;
		cout << median() << endl;
	}
}