[백준] 2696. 중앙값 구하기

문제 링크

[백준] 2696. 중앙값 구하기

풀이 과정

아이디어는 다음과 같습니다.

1. Min-heap 과 Max-heap 으로 수열을 관리한다.
2. Min-heap 에는 중앙값보다 큰 숫자들이 오름차순으로 정렬되고, Max-heap 에는 중앙값을 포함한 작은 숫자들이 내림차순으로 정렬된다.
3. Max-heap 의 요소 개수를 Min-heap 보다 한 개 많거나 같도록 값을 조정하여 유지한다.

다음은 예제를 시뮬레이션 해본 과정입니다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
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));
        StringTokenizer st;
        StringBuilder answer = new StringBuilder();
        int TC = Integer.parseInt(br.readLine());

        while (TC-- > 0) {
            int M = Integer.parseInt(br.readLine());
            int inputIdx = 0, outputIdx = 0, cnt = 1;
            PriorityQueue<Integer> minHeap = new PriorityQueue<>((a, b) -> a - b);
            PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);

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

            while (true) {
                if (inputIdx > M) break;

                String input = br.readLine();
                st = new StringTokenizer(input);

                while (st.hasMoreTokens()) {
                    int num = Integer.parseInt(st.nextToken());

                    if (minHeap.isEmpty() && maxHeap.isEmpty()) {
                        maxHeap.add(num);
                    } else if (maxHeap.size() > minHeap.size()) {
                        if (num > maxHeap.peek()) {
                            minHeap.add(num);
                        } else {
                            minHeap.add(maxHeap.poll());
                            maxHeap.add(num);
                        }
                    } else {
                        if (num > maxHeap.peek()) {
                            minHeap.add(num);
                            maxHeap.add(minHeap.poll());
                        } else {
                            maxHeap.add(num);
                        }
                    }

                    if (cnt % 2 == 1) {
                        outputIdx++;
                        answer.append(maxHeap.peek() + (outputIdx % 10 == 0 ? "\n" : " "));
                    }
                    cnt++;
                }

                inputIdx += 10;
            }
            answer.append("\n");
        }

        System.out.println(answer);
    }
}