[BOJ] 백준 1572번 중앙값 (Java)

#1572 중앙값

난이도 : 플레 5

유형 : 세그먼트 트리

1572번: 중앙값

▸ 문제

중앙값이란, 수열을 정렬했고, 그 크기가 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이다.

▸ 출력

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

 

문제 풀이  

세그먼트 트리와 슬라이딩 윈도우 기법을 활용하여 풀이를 해야 한다. 중앙값은 길이가 K인 부분 수열에서 (K+1)/2번째 수를 나타낸다. 모든 값을 한번에 삽입하여 구간을 반복하기보다는 슬라이딩 윈도우를 사용하면 된다. 즉, 이전 값은 버리고 다음 값을 추가하면서 길이가 K인 수열을 유지하면서 가운데 값을 계속 뽑아주는 식으로 답을 구한다.

 

수열 값 범위는 0부터 65536이기 때문에, 트리의 크기는 최대로 4*(65537)을 할당해주고 0부터 65536구간을 탐색해주면 된다. 왼쪽은 자신보다 작은 값, 오른쪽은 자신보다 큰 값의 개수를 나타내므로 다음과 같이 중앙값을 찾아주면 된다.

• tree[node*2] >= k,  왼쪽 자식 트리가 k보다 크므로 자신보다 작은 값이 나올 때까지 계속 재귀탐색한다.
• tree[node*2] < k, 왼쪽 자식 트리가 k보다 작으므로 자신보다 작은 개수를 빼주고 (cnt-tree[node*2]) 오른쪽 트리를 재귀 탐색한다.

 

만약 길이가 k=3인 [3, 4, 5]의 중앙값(k+1/2)을 구한다고 하면 다음과 같이 탐색한다. 

 

세그먼트 트리 중앙값

 

풀이 코드 

import java.io.*;
import java.util.*;

public class Main {
	
	static int[] arr, tree;
	static final int MAX = 65537;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		int n = Integer.parseInt(st.nextToken());
		int k = Integer.parseInt(st.nextToken());
		
		tree = new int[MAX*4];
		arr = new int[n+1];
		long sum = 0;
		for(int i=1; i<=n; i++) {
			arr[i] = Integer.parseInt(br.readLine());
			update(0, MAX, 1, arr[i], 1);
			if(i>=k) {
				int x = query(0,MAX,1,(k+1)/2);
				sum += x;
				update(0, MAX, 1, arr[i-k+1], -1);
			}
		}
		System.out.println(sum);
	}
	
	static void update(int s, int e, int node, int idx, int val) {
		if(idx < s|| e < idx) return;
		tree[node] += val;
		
		if(s == e) return;
		
		int mid = (s+e)/2;
		update(s, mid, node*2, idx, val);
		update(mid+1, e, node*2+1, idx, val);
	}
	
	static int query(int s, int e, int node ,int cnt) {
		if(s == e) return s;
		int mid = (s+e)/2;
		if(tree[node*2] >= cnt) return query(s, mid, node*2, cnt);
		else return query(mid+1, e, node*2+1, cnt-tree[node*2]);
	}
}