[ BOJ ] 1572 : 중앙값 ( PLATINUM 5 ) / C

문제

중앙값이란, 수열을 정렬했고, 그 크기가 N일 때, 1부터 시작해서 (N+1)/2번째 있는 원소가 그 수열의 중앙값이다. 예를 들어, {1, 2, 6, 5, 4, 3}에서는 3이고, {11, 13, 12, 15, 14}에서는 13이다.

오세준은 1초에 온도를 하나씩 재는 온도계를 만들었다. 이 온도계에는 작은 디스플레이 창이 하나 있는데, 이 창에는 지금부터 최근 K초 까지 온도의 중앙값을 표시해 준다. (온도를 재기시작한지 K초부터 표시한다. 그 전에는 아무것도 출력되지 않는다.)

오세준은 온도를 N초동안 쟀다. 그 시간 동안 온도계의 디스플레이 창에 뜨는 숫자의 합을 구하는 프로그램을 작성하시오.

다른 말로 하면, 길이가 N인 수열이 주어진다. 이 수열은 N-K+1 개의 길이가 K인 연속된 부분 수열이 존재한다. 이 부분 수열의 중앙값의 합을 출력하는 프로그램을 작성하시오.

입력

첫째 줄에 N과 K가 주어진다. N은 250,000보다 작거나 같은 자연수이고, K는 5,000보다 작거나 같은 자연수이다. N은 항상 K보다 크거나 같다. 둘째 줄부터 N개의 수가 한 줄에 하나씩 주어진다. 입력으로 주어지는 수는 65536보다 작거나 같은 자연수 또는 0이다.

출력

첫째 줄에 정답을 출력한다. 정답은 263-1보다 작거나 같다.

풀이 과정

슬라이딩 윈도우 기법을 이용했다. k 구간 만큼을 먼저 받고, "수열에 있는 값" 위치에 세그먼트 트리를 1 증가시켜 (k+1)/2번째 작은 수를 이분 탐색으로 찾을 수 있도록 설정했다. 다음 수열 값을 입력받을 때, "다음 수열 값" 위치에 세그먼트 트리를 1 증가시키고 가장 먼저 세그먼트 트리에 1을 증가시킨 값을 1 감소시켜 k 구간 만큼을 탐색할 수 있도록 하였다. 

#include <stdio.h>
#define SIZE 65550
typedef long long LL;

int arr[250001] = { 0, };
LL tree[SIZE * 4] = { 0, };

void update(int start, int end, int idx, int what, int value) {
	if (what < start || end < what) return;
	if (start == end) {
		tree[idx] += value;
		return;
	}

	int mid = (start + end) / 2;
	update(start, mid, idx * 2, what, value);
	update(mid + 1, end, idx * 2 + 1, what, value);
	tree[idx] = tree[idx * 2] + tree[idx * 2 + 1];
}

int findNum(int start, int end, int idx, int what) {
	if (start == end) return start;
	int mid = (start + end) / 2;

	if (tree[idx * 2] < what) return findNum(mid + 1, end, idx * 2 + 1, what - tree[idx * 2]);
	else return findNum(start, mid, idx * 2, what);
}

int main(void) {
	int n, k;
	scanf("%d %d", &n, &k);

	int i;
	int num;
	for (i = 0; i < k; i++) {
		scanf("%d", &num);
		arr[i] = num;
		update(0, SIZE - 1, 1, num + 1, 1);
	}

	LL res = 0;
	int tmp = 0;

	for (i = k; i < n; i++) {
		tmp = findNum(0, SIZE - 1, 1, (k + 1) / 2);
		res += (tmp - 1);

		scanf("%d", &num);
		arr[i] = num;

		update(0, SIZE - 1, 1, num + 1, 1);
		update(0, SIZE - 1, 1, arr[i - k] + 1, -1);
	}

	tmp = findNum(0, SIZE - 1, 1, (k + 1) / 2);
	res += (tmp - 1);

	printf("%lld", res);

	return 0;
}