[BOJ] 2696번 중앙값 구하기 - JAVA

💡 문제

💬 입출력 예시

📌 풀이(소스코드)

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.Collections;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

/**
 * 1시간
 * 다시 풀어야함. 기존 풀이는 두개의 리스트를 사용했는데, 두 개의 우선순위 큐를 사용하는 것이 훨씬 더 효율적임
 * 최대힙, 최소힙 두개를 사용해야함
 */
public class Main {

    public static void main(String[] args) throws Exception {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();

        int T = Integer.parseInt(br.readLine());
        while (T-- > 0) {
            // 최대힙과 최소힙 두개 사용. 최대힙은 내림차순 정렬, 최소힙은 오름차순 정렬 => 중앙값 접근
            PriorityQueue<Integer> minHeap = new PriorityQueue<>();
            PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

            // 출력하는 중앙값의 개수 == (수열의 크기 / 2) + 1
            int M = Integer.parseInt(br.readLine());
            sb.append((M + 1) / 2).append("\n");

            int cnt = 0;
            StringTokenizer st = null;
            for (int i = 0; i < M; i++) {
                // 입력을 위한 조건문
                if (i % 10 == 0) {
                    st = new StringTokenizer(br.readLine());
                }

                // 두 힙의 크기가 같으면 홀수번째 숫자 -> 최대힙에 저장
                // 두 힙의 크기가 다르면 짝수번째 숫자 -> 최소힙에 저장
                int num = Integer.parseInt(st.nextToken());
                if (maxHeap.size() == minHeap.size()) {
                    maxHeap.offer(num);
                } else {
                    minHeap.offer(num);
                }

                // 입력마다 최소 힙의 최솟값과 최대 힙의 최댓값을 비교
                // 최대 힙은 중앙 값 이하의 숫자만 가져야 한다. 즉, 최대 힙의 모든 값은 최소 힙의 모든 값보다 작아야 함
                // 그러므로 최대 힙의 최댓값과 최소 힙의 최솟값을 스왑해준다.
                if (!minHeap.isEmpty() && !maxHeap.isEmpty()) {
                    if (maxHeap.peek() > minHeap.peek()) {
                        int temp = minHeap.poll();
                        minHeap.offer(maxHeap.poll());
                        maxHeap.offer(temp);
                    }
                }

                // 홀수번째 수를 읽을 때마다 중앙값 출력
                // 위에서 최대 힙을 기준으로 하여 값을 추가했으므로 최대 힙의 최솟값이 중앙값
                if (i % 2 == 0) {
                    sb.append(maxHeap.peek()).append(" ");

                    // 한 줄에 10개씩 출력하기 위한 조건문
                    cnt++;
                    if (cnt % 10 == 0) {
                        sb.append("\n");
                    }
                }
            }
        }

        System.out.println(sb);
    }
}

📄 해설

접근

• 이전에 풀었던 풀이는 두 개의 ArrayList를 사용해서 중앙값을 저장하는 리스트와 수열을 저장하는 리스트 두개로 나누어서 풀었던 풀이이다.
• 그러나 이번에는 두 개의 우선순위 큐를 사용하는 풀이를 찾아 해당 풀이로 정리를 해보고자 한다.
• 최대 힙과 최소 힙 두개를 사용하는 것으로 접근을 해야하고, 어떤 것을 기준으로 하여 값을 추가할 것인지에 따라 출력부의 처리가 달라진다.

과정

• 이번 풀이 과정은 주석을 확인하는 것이 더 효율적일 것이라고 생각된다. 하지만 작성해보도록 하겠다. 입력과 출력을 위한 코드에 대한 과정은 알고리즘적인 부분은 아니므로 제외하겠다.
• 최소 힙과 최대 힙의 크기가 동일한지에 대해 확인한다. 이때 같으면 홀수번째 숫자가 들어온 것이므로 최대 힙에 넣고, 같지 않으면 짝수번째 숫자가 들어온 것이므로 최소 힙에 넣는다.
• 둘 중 하나의 힙에 값을 넣었다면, 두 힙 모두 비어있지 않을 경우 최소 힙의 최솟값인 minHeap.peek() 과 최대 힙의 최댓값인 maxHeap.peek() 을 비교한다. 최소 힙의 최솟값이 최대 힙의 최댓값보다 작은 경우, 최소 힙의 최댓값과 최대 힙의 최솟값을 스왑한다.
  • 이렇게 하는 이유는, 최소 힙의 모든 원소는 최대 힙의 모든 원소보다 커야하기 때문이고, 이를 유지하기 위함이다.
• 위 규칙을 그대로 따르고, 홀수번째 숫자의 입력 시 마다 출력문을 양식에 맞춰서 조정해주면 된다. 중앙값은 최대 힙을 기준으로 값을 입력했기 때문에, 최대 힙의 맨 처음 원소. 즉, 최대 힙의 최댓값이 된다.