[백준 1655번][Java] 가운데를 말해요

www.acmicpc.net/problem/1655

1655번: 가운데를 말해요

 

단순하게 접근하기

문제에서 N개의 다양한 크기의 정수가 주어집니다.

이 정수들이 하나씩 주어질 때마다 이전까지의 값들을 고려한 중간값을 찾아내야 합니다.

가장 단순한 접근은 값들을 리스트에 넣고 값이 주어질 때마다 리스트를 복사해 정렬한 뒤 중간값을 찾아내는 것입니다.

이럴 경우 리스트 복사 비용을 제외하더라도 시간 복잡도는 O(N^2 * logN)이 되기 때문에 시간초과로 실패합니다.

(N <= 100000)

---

중간값에 대해서 생각해보자

아래와 같이 정렬된 수열이 있다고 합시다.

1  2  3  4  5  6

중간값은 왼쪽 부분의 최대값, 오른쪽 부분의 최소값 둘 중 하나입니다.

이 경우엔 3이 더 작은 값이므로 3이 되겠네요.

 

이와 같이 값들을 작은 쪽, 큰 쪽으로 양분할 수 있다면

작은 쪽의 최대값, 큰 쪽의 최소값만 비교해 중간값을 찾아낼 수 있습니다.

최대, 최소값을 뽑아내기 위해서는 logN의 복잡도로 수행할 수 있는 우선순위 큐를 사용하면 되겠네요.

 

이제부터 왼쪽 회색 부분을 최대힙, 오른쪽 초록색 부분을 최소힙이라고 부르겠습니다.

 

---

정렬되지 않은 수에 대해서

정렬되지 않은 수열에 어떻게 위의 아이디어를 적용할 수 있을까요?

적절하게 값들을 옮겨주면 가능합니다. 

로직은 아래와 같습니다.

 

수를 힙에 추가하는 로직

1. 최대힙, 최소힙이 모두 비어있는 경우, 최대힙에 넣는다.

2. 최대힙의 최대값보다 주어진 수가 클 경우, 최소힙에 넣는다.

    최대힙의 최대값보다 주어진 수가 작거나 같을 경우, 최대힙에 넣는다.

3. 최대힙과 최소힙의 사이즈가 2개 이상 차이날 경우, 하나를 빼서 반대편 힙에 넣어준다.

 

중간값을 찾아내는 로직

1. 전체 사이즈가 홀수인 경우, 힙의 사이즈가 큰 쪽의 루트가 중간값이다.

2. 짝수라면 각 힙의 루트 중 작은 값이 중간값이다. 

 

---

예제

 

아래와 같은 수열이 있다고 합시다.

 

5  3  2  1  6  4

 

- 처음에 5가 주어집니다.

최대힙, 최소힙 모두 비어있으므로 일단 최대합에 5를 넣어줍니다.

 

최대힙 : [5]

최소힙 : []

 

중간값 : 5  

---

- 3이 주어집니다.

 5 > 3이므로 3이 최대힙에 들어가고, 5를 빼내 최소힙에 넣습니다.

최대힙 : [3] 

최소힙 : [5]

 

중간값 : 3

---

- 2가 주어집니다.

3 > 2이므로 최대힙에 들어가지만, 크기가 2개이상 차이나지 않기 때문에 이동은 없습니다.

 

최대힙 : [3, 2]

최소힙 : [5]

 

중간값 : 3

---

- 1이 주어집니다.

3 > 1이므로 1이 최대힙에 들어가고, 최대힙의 루트를 빼내 최소힙에 넣습니다.

최대힙 : [2, 1]

최소힙 : [3, 5]

 

중간값 : 2

---

- 6이 주어집니다.

2 < 6이므로 최소힙에 넣어줍니다.

 

최대힙 : [2, 1]

최소힙 : [3, 5, 6]

 

중간값 : 3 (개수가 홀수일 경우, 중간값은 반드시 힙 사이즈가 큰 쪽에 있습니다.)

---

- 4가 주어집니다.

2 < 4이므로 최소힙에 넣어주고, 최소힙의 루트를 최대힙에 넣어줍니다.

 

최대힙 : [3, 2, 1]

최소힙 : [4, 5, 6]

 

중간값 : 3

---

위 예제처럼 최대힙의 루트를 기준으로 값을 각 힙에 넣어주고, 개수 차이가 2개 이상 날때마다 값을 다시 옮겨주면

크기 기준으로 값들을 양분하는 두 개의 힙을 만들 수 있습니다.

 

그러고나면 중간값은 이 두 힙의 루트를 비교해 간단하게 얻을 수 있지요.

개수가 홀수인 경우엔 크기가 큰 힙의 루트가 중간값이 되고, 짝수인 경우엔 힙들의 루트 중 작은값이 중간값이 됩니다. 

 

 

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Comparator;
import java.util.PriorityQueue;

public class Main {
    static int N;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
//        BufferedReader br = new BufferedReader(new FileReader("testcase.txt"));
        N = Integer.parseInt(br.readLine());

        PriorityQueue<Integer> minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a));
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.comparingInt(a -> -a));

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < N; i++) {
            int inputNumber = Integer.parseInt(br.readLine());
            add(inputNumber, minHeap, maxHeap);
            int midVal = getMid(minHeap, maxHeap);
            sb.append(midVal);
            if(i != N-1) sb.append("\n");
        }

        System.out.println(sb);
    }

    private static void add(int inputNumber, PriorityQueue<Integer> minHeap, PriorityQueue<Integer> maxHeap) {
        if(minHeap.isEmpty() && maxHeap.isEmpty()) {
            maxHeap.add(inputNumber); return;
        }

        if(maxHeap.peek() < inputNumber) {
            minHeap.add(inputNumber);
        }else {
            maxHeap.add(inputNumber);
        }

        if(maxHeap.size() - minHeap.size() == 2) {
            int maxVal = maxHeap.poll();
            minHeap.add(maxVal);
        }

        if(minHeap.size() - maxHeap.size() == 2) {
            int minVal = minHeap.poll();
            maxHeap.add(minVal);
        }
    }

    private static int getMid(PriorityQueue<Integer> minHeap, PriorityQueue<Integer> maxHeap) {
        if(minHeap.size() == maxHeap.size()) {
            return Math.min(minHeap.peek(), maxHeap.peek());
        }

        /*
        주어진 입력의 개수가 홀수인 경우, heap 사이즈가 하나 큰 쪽이 답을 가진다.
         */
        if(minHeap.size() > maxHeap.size()) {
            return minHeap.peek();
        }
        return maxHeap.peek();

    }
}