[Java] 백준 2696번: 중앙값 구하기 - Priority Queue(우선순위 큐)

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

2696번: 중앙값 구하기

문제

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

예를 들어, 수열이 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개씩 출력해야 한다.

 

---

 

문제를 보고 처음 든 생각은 홀수번째 수일 때 마다 정렬을 해서 중앙값을 가져오는 것..

너무 비효율적일 것 같다는 생각에 다른 블로그를 참고하여 우선순위 큐로 풀었다!

 

최대힙과 최소힙으로 나누어서 각각의 값들을 뽑아서 비교하며 위치를 정한다.

 

최대힙은 최대값을 출력하도록 내림차순 정렬

최소힙은 최소값을 출력하도록 오름차순 정렬

PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
PriorityQueue<Integer> minHeap = new PriorityQueue<>();

 

번갈아가면서 각각의 힙에 값을 넣어준다.

if(maxHeap.size() == minHeap.size())
    maxHeap.offer(x);
else
    minHeap.offer(x);

 

힙에 값이 들어왔으니 최소힙과 최대힙의 값을 꺼내서 비교한다.

최대힙의 최대값 > 최소힙의 최소값 이라면 두 값을 교환한다.

if(maxHeap.peek() > minHeap.peek()){
    int t1 = maxHeap.poll();
    int t2 = minHeap.poll();

    maxHeap.offer(t2);
    minHeap.offer(t1);
}

 

그러면 최대힙에는 최소힙에 있는 값들보다 작은 값들만 들어가게 된다.

(즉, 중앙값보다 같거나 작은 값들이 들어감)

교환이기때문에 두 힙의 개수에는 변동이 없고,

 작은 값, 큰 값이 상대적으로 비교되어 들어가므로 중앙값을 구할 수 있다.

 

처음에는 "제일 큰 값을 최소 힙에 넣었다고 해도 최대힙에 아직 큰 값이 남아있을 수 도 있지 않나" 라는 생각을 하였다.
→ 값을 넣어줄 때마다 비교하므로 남아있을 수 가 없다!

 

 

마지막으로, 홀수번째일 때만 최대힙에서 중앙값을 가져오면 된다.

배열의 index는 0부터 시작하므로 짝수번째일 경우에 가져오면 된다.

if(i%2 == 0){
    if(cnt == 9 || i == M - 1){
        sb.append(maxHeap.peek() + "\n");
        cnt = 0;
    } else{
        sb.append(maxHeap.peek() + " ");
    }
    cnt++;
}

 

소스 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Main {
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringBuilder sb = new StringBuilder();
		StringTokenizer st = null;
		int T = Integer.parseInt(br.readLine());
		for (int t = 0; t < T; t++) {
			PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
			PriorityQueue<Integer> minHeap = new PriorityQueue<>();
			int M = Integer.parseInt(br.readLine());

			sb.append(M/2+1 + "\n");

			int cnt = 0;

			for (int i = 0; i < M; i++) {
				if (i % 10 == 0) {
					st = new StringTokenizer(br.readLine());
				}

				int x = Integer.parseInt(st.nextToken());

				if(maxHeap.size() == minHeap.size())
					maxHeap.offer(x);
				else
					minHeap.offer(x);

				if(!minHeap.isEmpty()){
					if(maxHeap.peek() > minHeap.peek()){
						int t1 = maxHeap.poll();
						int t2 = minHeap.poll();

						maxHeap.offer(t2);
						minHeap.offer(t1);
					}
				}

				if(i%2 == 0){
					if(cnt == 9 || i == M - 1){
						sb.append(maxHeap.peek() + "\n");
						cnt = 0;
					} else{
						sb.append(maxHeap.peek() + " ");
					}
					cnt++;
				}
			}
		}

		System.out.println(sb);
	}
}