[Boj_2696] 중앙값 구하기

📎 문제 링크

https://www.acmicpc.net/problem/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개씩 출력해야 한다.

 

📍 문제 정보

• 🥇  문제 레벨 : 골드 2
• 🔔  문제 유형 : 자료 구조, 우선순위 큐 
• 💬  풀이 언어 : JAVA

 

🤔 문제 풀이

백준 1655 가운데를 말해요와 비슷한 문제이다.

 

우선순위 큐 두 개를 사용하여 중앙값을 찾을 수 있다.

 

먼저, 최대 힙과 최소 힙을 구현할 우선순위 큐 두 개를 선언한다.

최대 힙은 최댓값을 정렬하도록 내림차순으로 정렬하고,

최소 힙은 최솟값을 정렬하도록 오름차순으로 정렬한다.

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

 

번갈아가며 각각 힙에 값을 넣어준다.

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

 

최대 힙의 최댓값 > 최소 힙의 최솟값일때, 두 값을 교환한다.

if (maxHeap.peek() > minHeap.peek()) {
    int tmp = minHeap.poll();
    minHeap.offer(maxHeap.poll());
    maxHeap.offer(tmp);
}

 

그러면 최대 힙에는 중앙값 이하의 수만, 최소 힙에는 중앙값 초과의 수만 존재하게 된다.

그렇기 때문에, 최대 힙에는 항상 중앙값이 존재하게 된다.

 

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

 

위 풀이를 코드로 작성해 보면 아래와 같다.

 

✔ 최종 코드

메모리 : 15924KB
시간 : 152ms

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()); // 테스트 케이스 수
        while (t --> 0) {
            // 최대힙
            PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
            // 최소힙
            PriorityQueue<Integer> minHeap = new PriorityQueue<>();

            int m = Integer.parseInt(br.readLine());

            sb.append((m/2)+1).append("\n"); // 중앙 값 개수

            int cnt = 0;
            for (int i = 0; i < m; i++) {
                if (i % 10 == 0) { // 원소는 한줄에 10개씩
                    st = new StringTokenizer(br.readLine());
                }
                
                // 중앙 값 구하기
                int num = Integer.parseInt(st.nextToken());
                if (minHeap.size() == maxHeap.size()) maxHeap.offer(num);
                else minHeap.offer(num);

                if (!minHeap.isEmpty() && !maxHeap.isEmpty()) {
                    if (maxHeap.peek() > minHeap.peek()) {
                        int tmp = minHeap.poll();
                        minHeap.offer(maxHeap.poll());
                        maxHeap.offer(tmp);
                    }
                }

                // 홀수 번째 수를 읽을 때마다
                if (i % 2 == 0) {
                    // 출력도 한 줄에 원소 10개씩
                    if (cnt == 9 || i == m-1) {
                        sb.append(maxHeap.peek()).append("\n");
                        cnt = 0;
                    } else {
                        sb.append(maxHeap.peek()).append(" ");
                    }
                    cnt++;
                }
            }
        }

        System.out.println(sb);
    }
}

 

전에 가운데를 말해요 문제를 풀면서 우선순위 큐 두 개를 사용하여 중앙값을 찾는 다른 사람의 풀이가 굉장히 신기하고 인상 깊었다. 그래서 이번 문제를 풀 때 그때 풀이를 활용했고, 블로그에 포스팅하며 다시 한번 정리했다.

부지런히 더 열심히 공부해야지...💭