[백준] 1655번 : 가운데를 말해요 - JAVA [자바]

문제 링크

문제

백준이는 동생에게 "가운데를 말해요" 게임을 가르쳐주고 있다. 백준이가 정수를 하나씩 외칠때마다 동생은 지금까지 백준이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 백준이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 백준이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 백준이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에는 백준이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 백준이가 외치는 정수가 차례대로 주어진다. 정수는 -10,000보다 크거나 같고, 10,000보다 작거나 같다.

출력

한 줄에 하나씩 N줄에 걸쳐 백준이의 동생이 말해야 하는 수를 순서대로 출력한다.

풀이

당연하게도, 매 입력마다 값을 정렬하여 중간 인덱스를 통하여 중간값을 찾게되면 시간 초과가 나겠다는 생각을 했다.

이에 대한 대안으로 입력 받은 숫자를 왼쪽 부분과 오른쪽 부분 나누어 저장하면 되겠다는 생각을 했다.

왼쪽 부분은 maxHeap, 오른쪽 부분은 minHeap으로 설정하여 새로운 값이 현재의 중간값보다 크면 최소힙(오른쪽)에, 작으면 최대 힙(왼쪽)에 넣으면 항상 중간값을 maxHeap에 유지할 수 있겠다 생각했다.

새로운 값이 입력되고난 후에 두 Heap의 크기를 1 이하로 유지하는 것도 빼먹으면 안됐다.

import java.io.*;  
import java.util.*;  
  
public class Main {  
  
    public static void main(String[] args) throws IOException {  
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));  
        int n = Integer.parseInt(br.readLine());  
        StringBuilder sb = new StringBuilder();  
  
        // 최소 힙과 최대 힙 생성  
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();  
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());  
  
        for (int i = 0; i < n; i++) {  
            int x = Integer.parseInt(br.readLine());  
  
            // 최대 힙이 비어 있거나 x가 최대 힙의 루트보다 작거나 같으면, 최대 힙에 x를 추가  
            if (maxHeap.isEmpty() || x <= maxHeap.peek()) {  
                maxHeap.offer(x);  
            } else { // x가 최대 힙의 루트보다 크면 최소 힙에 x를 추가  
                minHeap.offer(x);  
            }  
  
            // 최대 힙의 크기가 최소 힙의 크기 + 1보다 클 경우, 최대 힙의 루트를 최소 힙으로 이동  
            if (maxHeap.size() > minHeap.size() + 1) {  
                minHeap.offer(maxHeap.poll());  
            } else if (minHeap.size() > maxHeap.size()) {  
                // 최소 힙의 크기가 최대 힙의 크기보다 크면, 최소 힙의 루트를 최대 힙으로 이동  
                maxHeap.offer(minHeap.poll());  
            }  
  
            // 중간값 출력  
            sb.append(maxHeap.peek()).append("\n");  
        }  
        System.out.println(sb);  
    }  
}

주어진 입력을 통해 위 코드는

[1] []
[1] [5]
[1, 2] [5]
[1, 2] [5, 10]
[-99, 1, 2] [5, 10]
[-99, 1, 2] [5, 7, 10]
[-99, 1, 2, 5] [5, 7, 10]

위와 같이 동작하게 된다. 즉 항상 왼쪽의 maxHeap의 값을 출력하면 문제에서 원하는 중간값을 출력할 수 있다.

PriorityQueue으로 최대힙 구현하기(maxHeap)

지난번 문제였던 최대 힙 문제에서

PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparingInt(x -> -x));

로 최대 힙을 구현했는데, 잘 찾아보니..

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

위와 같이. Comparator.reverseOrder()를 인자로 넘기면 되었다.