[SWEA] 3000. 중간값 구하기 (with JAVA)

문제 3000

SW Expert Academy

 

중간값을 구하는 방법 중에서 우선순위 큐 (Priority Queue)를 사용해보겠습니다. 삽입할 때마다 중간값을 구하는 문제로, Heap 으로 구성된 Priority Queue 를 사용합니다. 시간 복잡도는 O(nlogn) 입니다.

중간값을 구하기 위해서는 아래와 같은 규칙을 따라야 합니다.
1. 제일 처음에는 최대 힙에 삽입
2. 최대 힙의 크기는 최소 힙의 크기와 같거나 하나 더 큼
3. 최대 힙의 최대 원소는 최소 힙의 최소 원소보다 작거나 같음. 즉, 최대 힙의 루트 노드 <= 최소 힙의 루트 노드
    - 알고리즘에 맞지 않다면 최대 힙, 최소 힙의 가장 위의 값을 swap
4. 이 규칙을 유지해준다면 항상 최대 힙의 루트 노드 값이 중간값

위의 규칙을 적용하여 중간값을 구해보겠습니다. 이 문제에서는 맨 처음에 5를 삽입하고, 그 후에 [1, 3], [2, 6], [8, 9] 를 차례대로 삽입합니다. 이 과정을 그림으로 나타내면 다음과 같습니다.

1) 1번 규칙에 따라, 맨 처음에 5를 최대 힙에 삽입합니다.

최대 힙에 5 삽입

 

2) 다음 숫자 [1, 3] 을 최소, 최대 힙에 순서대로 삽입합니다. (** 최소 힙에 먼저 삽입!)

최소, 최대 힙 순으로 삽입

 

3) 2번 규칙에 따라서 최대 힙의 루트 노드가 더 작아야 하기 때문에 swap 해줍니다. 이 때, 중간값은 3 입니다.

swap

 

4) 이제 다음 숫자인 [2, 6] 을 위와 같은 방법으로 삽입해주면 아래 그림과 같습니다.

[2, 6] 삽입

 

5) 2번 규칙에 따라 최대 힙 루트 노드가 더 작아야 하기 때문에 swap 해줍니다. 이 때, 중간값은 3 입니다.

swap

 

6) 마지막으로 [8, 9] 를 삽입해줍니다.

[8, 9] 삽입

 

7) 최대 힙 루트 노드와 최소 힙 루트 노드를 비교해서 swap 합니다. 이 때, 중간값은 5입니다.

swap

위의 내용을 바탕으로 구현해보면 아래와 같습니다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
    private static PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    private static PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = Integer.parseInt(br.readLine());
        StringTokenizer st;
        StringBuilder sb = new StringBuilder();

        for (int test_case = 1; test_case <= T; test_case++) {
            minHeap.clear();
            maxHeap.clear();
            st = new StringTokenizer(br.readLine());
            int N = Integer.parseInt(st.nextToken());
            int A = Integer.parseInt(st.nextToken());
            maxHeap.offer(A); // 최대 힙에 먼저 삽입

            long res = 0;
            for (int i=0; i<N; i++) {
                st = new StringTokenizer(br.readLine());
                minHeap.offer(Integer.parseInt(st.nextToken()));
                maxHeap.offer(Integer.parseInt(st.nextToken()));

                // 최대힙 루트노드와 최소힙 루트노드 비교해서 swap (최대힙 루트노드가 더 작아야 함)
                if (!maxHeap.isEmpty() && !minHeap.isEmpty()) {
                    if (maxHeap.peek() > minHeap.peek()) {
                        int tmp = maxHeap.poll();
                        maxHeap.offer(minHeap.poll());
                        minHeap.offer(tmp);
                    }
                    // 최대힙 루트노드 값이 중간값
                    res += maxHeap.peek();
                }
            }
            sb.append("#").append(test_case).append(" ").append(res%20171109).append("\n");
        }
        System.out.println(sb);
    }
}