[BOJ] 백준 2696번 : 중앙값 구하기 (JAVA)

문제

어떤 수열을 읽고, 홀수번째 수를 읽을 때 마다, 지금까지 입력받은 값의 중앙값을 출력하는 프로그램을 작성하시오.

예를 들어, 수열이 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비트 부호있는 정수이다. (대부분의 언어에서 int)

 

 

출력

각 테스트 케이스에 대해 첫째 줄에 출력하는 중앙값의 개수를 출력하고, 둘째 줄에는 홀수 번째 수를 읽을 때 마다 구한 중앙값을 차례대로 공백으로 구분하여 출력한다. 이때, 한 줄에 10개씩 출력해야 한다.

 

 

풀이

무식하게 푸는 방법과 우선순위 큐를 이용한 방식 2가지로 풀 수 있었습니다. 이와는 별도로 출력에 있어서 주의할 점이 있었는데, 한 줄의 최대 10개까지만 출력할 수 있다는 점입니다. 따라서, 특정 변수를 두어서 제한된 출력 개수만큼 올바르게 출력하는 것이 중요했습니다.

 

 

단순한 정렬을 이용한 풀이

제가 가장 먼저 생각한 방식은 입력받을 때마다 배열을 정렬하여 중앙값을 출력하고자 하였습니다. 사실, 처음에는 시간 복잡도 상으로 시간초과가 날 것으로 예상하였으나, 테스트 케이스 상으로 운이 좋아서 통과한 것이 아닌가 생각하고 있습니다.

 

그래도 일단, 1차적으로 성공은 했으므로 간략하게 로직을 알려드리도록 하겠습니다.

 

1. N을 입력받고, 중앙값의 개수를 출력한다. -> (N + 1) / 2를 하면 됨.

2. 배열을 하나 정의하고, 값을 입력받을 때마다 오름차순으로 정렬한다.

3. 인덱스는 0부터 시작하므로 짝수 번째 인덱스를 접근할 때마다 배열의 (i / 2)번째 값을 출력한다.

4. cnt라는 변수를 정의하여, cnt = 9가 될 경우, 출력할 때 한 줄 띄도록 조정한다.

 

아래는 위 과정을 정리한 소스코드입니다.

 

 

소스코드

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;
 
public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = null;
        int T = Integer.parseInt(br.readLine()); // 테스트 케이스의 개수
 
        StringBuilder sb = new StringBuilder();
        while (T-- > 0) {
            int N = Integer.parseInt(br.readLine()); // 배열의 크기
 
            sb.append(((N + 1) / 2) + "\n"); // 중앙값의 개수 출력
            ArrayList<Integer> a = new ArrayList<>();
            int cnt = 0; // 줄 간격 띄우기 위한 용도
 
            for (int i = 0; i < N; i++) {
                if (i % 10 == 0) {
                    st = new StringTokenizer(br.readLine());
                }
 
                int x = Integer.parseInt(st.nextToken());
                a.add(x);
                Collections.sort(a); // 입력받을 때마다 오름차순 정렬.
                
                // 인덱스는 0부터 시작하므로 짝수 번째 인덱스를 탐색할 때마다
                // 중앙값을 출력하면 됨.
                if (i % 2 == 0) {
                    // 한 줄의 최대 10개만 가능.
                    if (cnt == 9 || i == N - 1) {
                        sb.append(a.get(i / 2) + "\n");
                        cnt = 0;
                    } else {
                        sb.append(a.get(i / 2) + " ");
                    }
                    cnt++;
                }
            }
        }
 
        bw.write(sb.toString());
        bw.flush();
        bw.close();
        br.close();
    }
 
}

 

 

 

우선순위 큐를 이용한 풀이

위의 풀이는 시간 초과의 위험이 있지만, 우선순위 큐는 매 순간 모든 요소를 정렬하여 중앙값을 구할 필요 없이, 바로 중앙값만을 가져오는 데 용이합니다.

 

핵심은 최대 힙과 최소 힙을 설정하는 것이었습니다. 최대 힙은 중앙값 이하의 수가 오도록, 최소 힙은 중앙값 초과의 수가 오도록 조정하면 됩니다. 이 때, 최대 힙은 내림차순 정렬, 최소 힙은 오름차순 정렬이 되도록 해야 합니다.

 

위를 기반으로 로직을 정리하면 아래와 같습니다.

 

1. N을 입력받고, 중앙값의 개수를 출력한다. -> (N + 1) / 2를 하면 됨.

2. 최대 힙과 최소 힙의 사이즈가 같을 경우, 입력받은 값을 최대 힙에 넣고, 사이즈가 다른 경우, 입력받은 값을 최소 힙에 넣는다. -> 왼쪽과 오른쪽 우선순위 큐에 차례대로 1개씩 넣는다는 느낌.

3. 최소 힙이 비어있지 않은 경우, 최대 힙의 peek()한 값과 최소 힙의 peek()한 값을 비교한다.

3-1. 만약, 최대 힙의 peek()한 값이 최소 힙의 peek()한 값보다 클 경우, 두 개의 값을 서로 바꾼다.

4. 인덱스는 0부터 시작하므로, 짝수 번째가 될 때마다 maxHeap의 peek()한 값을 출력한다.

5. cnt라는 변수를 정의하여, cnt = 9가 될 경우, 출력할 때 한 줄 띄도록 조정한다.

 

3-1 과정을 하는 이유는 단순합니다. maxHeap에는 "중앙값 이하의 수만 가져야 한다."는 조건을 충족시키기 위해서입니다. 

 

아래는 위 과정을 정리한 소스코드입니다.

 

 

소스코드

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
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));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = null;
        int T = Integer.parseInt(br.readLine()); // 테스트 케이스의 개수
 
        StringBuilder sb = new StringBuilder();
        while (T-- > 0) {
            PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
            PriorityQueue<Integer> minHeap = new PriorityQueue<>();
            int N = Integer.parseInt(br.readLine()); // 배열의 크기
 
            sb.append(((N + 1) / 2) + "\n"); // 중앙값의 개수 출력
 
            int cnt = 0; // 줄 간격 띄우기 위한 용도
 
            for (int i = 0; i < N; i++) {
                if (i % 10 == 0) {
                    st = new StringTokenizer(br.readLine());
                }
 
                int x = Integer.parseInt(st.nextToken());
 
                // 입력받는 값들을 하나씩 차례대로 왼쪽, 오른쪽에 넣는 느낌.
                if (maxHeap.size() == minHeap.size()) {
                    maxHeap.offer(x);
                } else {
                    minHeap.offer(x);
                }
 
                // maxHeap은 중앙값 이하의 숫자만 갖도록 조정.
                if (!minHeap.isEmpty()) {
                    if (maxHeap.peek() > minHeap.peek()) {
                        int t1 = maxHeap.poll();
                        int t2 = minHeap.poll();
 
                        maxHeap.offer(t2);
                        minHeap.offer(t1);
                    }
                }
 
                // 인덱스는 0부터 시작하므로 짝수 번째 인덱스를 탐색할 때마다
                // 중앙값을 출력하면 됨.
                if (i % 2 == 0) {
                    // 한 줄의 최대 10개만 가능.
                    if (cnt == 9 || i == N - 1) {
                        sb.append(maxHeap.peek() + "\n");
                        cnt = 0;
                    } else {
                        sb.append(maxHeap.peek() + " ");
                    }
                    cnt++;
                }
            }
        }
 
        bw.write(sb.toString());
        bw.flush();
        bw.close();
        br.close();
    }
 
}

 

지적 혹은 조언 환영합니다! 언제든지 댓글로 남겨주세요.