[Java] 최소 힙, 최대 힙를 사용하여 중앙 값 구하기

 

https://www.acmicpc.net/problem/2696

https://www.acmicpc.net/problem/1655

해당 문제의 경우 입력받은 숫자들 중에서 중앙값을 구하는 문제이다. 예를 들어, [1, 5, 7, 8, 9] 리스트가 존재할 때 중앙에 위치한 값은 7이다.

 

본 문제를 풀기 위해, 1개의 우선순위 큐를 사용하여 문제를 풀었다.

자바에서 PriorityQueue를 사용하여 원소 삽입, 삭제 시에 발생하는 시간 복잡도는 다음과 같다.

• 삽입 : O(log N)
• 삭제 : O(log N)

원소 삽입 및 삭제 시에 최소 힙, 최대 힙 성질을 만족하기 위해 트리 구조를 조정한다. 

최소 힙과 최대 힙 그리고 트리 구조 조정에 대해서 모른다면 다음 글을 참고한다. 

https://bristle-comb-69b.notion.site/14343369026b80c88691fff9a66cc724?pvs=4

 

아래 코드는 pq에 담긴 원소들 중에서 중간에 위치한 값을 추출하기 위해 getMiddleNumber()를 호출한다. 이 메서드는 다음과 같은 이유로 비효율적으로 동작한다.

1. n개가 담긴 원소들 중에서 절반(n/2)을 추출해야 한다. 원소를 추출할 때마다 최소, 최대 힙의 성질을 만족하기 위해 O(log N)의 시간 복잡도가 발생한다.
2. 중간 값을 추출한 뒤, 1번에서 추출한 n/2 개의 원소를 다시 삽입한다. 원소를 삽입할 때마다 최소, 최대 힙 성질을 만족하기 위해 O(log N)의 시간 복잡도가 발생한다.

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        StringBuilder sb = new StringBuilder();
        int T = Integer.parseInt(br.readLine());
        while (T-- > 0) {
            int n = Integer.parseInt(br.readLine()); // 원소 개수
            int m = n / 10; // 몇 줄 읽은 것인지?

            PriorityQueue<Integer> pq = new PriorityQueue<>(); // 우선순위 큐
            List<Integer> result = new ArrayList<>(); // 출력할 중간 값들
            boolean isOdd = true; // 홀수 여부, true : 홀수

            // 원소를 10개씩 끊어서 읽는다.
            // ex. 원소 개수가 23개일 경우, 3줄에 걸쳐서 읽는다.
            for (int i = 0; i <= m; i++) {
                String[] inputs = br.readLine().split(" ");
                for (int j = 0; j < inputs.length; j++) {
                    pq.offer(Integer.parseInt(inputs[j]));

                    // 홀수라면 중간 값을 result에 저장한다.
                    if (isOdd) {
                        result.add(getMiddleNumber(pq));
                    }
                    isOdd = !isOdd;
                }
            }

            // 중간 값들 출력
            sb.append(result.size()).append("\n");
            for (int i = 0; i < result.size(); i++) {
                sb.append(result.get(i));
                // 각 줄에 10개의 원소를 출력한다.
                // ex. 출력할 원소가 12개라면, 2줄에 걸쳐서 출력한다.
                sb.append((i + 1) % 10 == 0 ? "\n" : " ");
            }
            sb.append("\n");
        }

        System.out.println(sb);
    }

    static int getMiddleNumber(PriorityQueue<Integer> pq) {
        // 중간 값을 추출하기 위해, 임시로 추출한 값을 저장
        Stack<Integer> stack = new Stack<>();

        // 중간 값을 추출한다.
        int middleIdx = pq.size() / 2;
        for (int i = 0; i < middleIdx; i++) {
            stack.push(pq.poll());
        }
        int middleNum = pq.peek();

        // 추출한 원소들을 다시 큐에 저장한다.
        while (!stack.isEmpty()) {
            pq.offer(stack.pop());
        }

        // 중간 값 반환
        return middleNum;
    }
}

 

최대 9999개의 숫자가 입력될 수 있으므로, 홀수 번째마다 getMiddleNumber() 작업을 수행하는 것은 매우 비효율적이다. 실제로 아래 사진을 보면 수행 시간이 상당이 오래 걸린다.

 

 

 

 

두 개의 우선순위 큐를 사용하면 좀 더 효율적으로 코드를 작성할 수 있다. 최소 힙, 최대 힙을 사용하여 중앙값을 구하는 테크닉을 이용하면 된다. (기존에는 하나의 우선순위 큐만 사용했었다.)

 

최소 힙, 최대 힙을 사용하여 중앙값을 구하는 프로세스는 다음과 같다.

1. 최소 힙과 최대 힙에 저장된 원소의 개수가 동일한지 확인한다.
   1. 동일하다면, 최대 힙에 원소를 삽입한다.
   2. 동일하지 않다면, 최소 힙에 원소를 삽입한다.
2. 최소 힙이 비어있지 않다면
   1. 최소 힙의 root 원소와 최대 힙의 root 원소를 비교한다.
   2. 최대 힙 > 최소 힙이라면, 최소 힙과 최대 힙의 Root 원소를 교환한다.

 

위 과정을 수행하면 최대 힙에 중앙값이 저장된다. 

• 최대 힙 : [2, 1]
• 최소 힙 : [3]

최소, 최대 힙의 상태가 위와 같을 때, 최대 힙의 root 원소에 2가 위치한다. 이 값이 현재까지 입력받은 [1, 2, 3] 숫자 중에서 중앙값에 해당한다.

 

중앙값을 구하는 방법과 추출하는 방법을 코드로 작성하면 다음과 같다.

 

 

최소 힙과 최대 힙을 사용하여 https://www.acmicpc.net/problem/2696 문제를 해결한 코드이다.

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        // T번의 과정을 반복한다.
        int T = Integer.parseInt(br.readLine());

        StringBuilder sb = new StringBuilder();
        while (T-- > 0) {
            int n = Integer.parseInt(br.readLine()); // 원소 개수
            sb.append((n + 1) / 2).append("\n"); // 중앙 값의 개수

            // 두 개의 자료구조를 사용하여 중앙 값을 구한다.
            PriorityQueue<Integer> minHeap = new PriorityQueue<>();
            PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

            int m = n / 10; // 몇 줄 입력받을 것인지?
            int printCnt = 0; // 출력 횟수, 한 줄에는 최대 10개의 숫자만 출력될 수 있다.
            for (int i = 0; i <= m; i++) {
                String[] inputs = br.readLine().split(" ");
                for (int j = 0; j < inputs.length; j++) {
                    int val = Integer.parseInt(inputs[j]); // 입력받은 숫자

                    // 최소 힙, 최대 힙 크기에 따라 숫자를 삽입할 위치를 결정한다.
                    if (minHeap.size() == maxHeap.size()) { // 사이즈가 동일하다면
                        maxHeap.add(val);
                    } else {
                        minHeap.add(val);
                    }

                    // 최소 힙의 Root 원소 < 최대 힙의 Root 원소라면,
                    // 두 원소를 교체한다.
                    if (!minHeap.isEmpty()) {
                        if (maxHeap.peek() > minHeap.peek()) {
                            swap(minHeap, maxHeap);
                        }
                    }

                    // 홀수 번째의 원소일 때, 중앙 값을 출력한다.
                    if (j % 2 == 0) {
                        sb.append(maxHeap.peek()).append(" ");

                        // 한 줄에 최대 10개의 원소를 출력할 수 있다.
                        if (++printCnt % 10 == 0) {
                            sb.append("\n");
                        }
                    }
                }
            }
            sb.append("\n");
        }

        System.out.println(sb);
    }

    static void swap(PriorityQueue<Integer> minHeap, PriorityQueue<Integer> maxHeap) {
        // 최대 힙 Root > 최소 힙 Root이므로,
        // 최대 힙에 최소 힙의 Root를 삽입해도, 여전히 최대 힙의 Root는 동일하다.
        maxHeap.add(minHeap.poll());
        minHeap.add(maxHeap.poll());
    }
}

 

 

 

 

참고

https://m.blog.naver.com/adamdoha/222145579848