백준 7662번 : 이중 우선순위 큐에 대한 고찰 (Java, Python)

목차

서론

삽질의 결과.. 다음 페이지가 눈에 보인다면 당신의 눈은 정상입니다.

PS 세상에서 난 행복할 수 없는걸까... 

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

7662번: 이중 우선순위 큐

응 햄보칼수없어..

한달 전에 풀었었지만 오늘 프로그래머스에서 동일한 문제를 풀다가 생각이나서 Java로 다시 풀어봤다!

https://school.programmers.co.kr/learn/courses/30/lessons/42628

참고로 프로그래머스에 있는 문제와 백준에 올라와 있는 문제는 동일하다. 같은 로직을 써서 정답을 받을 수 있지만 프로그래머스의 채점 테스트케이스가 6개뿐이라서 비교적 쉽게 통과되니까 이왕이면 같은 로직을 백준에서 돌려보길 권장한다.

 

풀이

위 문제는 기존의 우선순위 큐와 다르게 최대 힙, 최소 힙 만들어서 푸는게 핵심이다. 따라서 최대 힙에서는 최대값만 제거하고 최소힙에서는 최소값을 제거하고 각각 제거하는 값들을 visited 배열에서 추적해주면 된다. 처음 힙에 들어갈때는 참으로 한번이라도 힙에서 제거되면 거짓으로 만들어 값을 추적한다.

 

모든 데이터 삭제와 삽입이 끝나면 visited 배열에서 존재하지 않아야 하는 값들이 있는지 확인한다. 이에 따라서 불필요한 값들을 각각 최대 힙과 최소 힙에서 지울 수 있다.

 

총 테스트케이스의 수가 T이고 각 테스트케이스에서 연산의 수 (최대 1,000,000) 만큼 반복문을 돌아야하니 시간복잡도는 O(T * 1000000) 이 나올 것이다. 

 

이외의 방법에서는 Java는 TreeMap (균형 이진 트리)를 통해 최대값과 최솟값을 동시에 빼고 넣을 수 있다. Python은 비슷한 자료구조가 표준 라이브러리에서는 없어서 이 방법을 통해 문제를 해결하려면 직접 구현 해야한다.

 

Java - 최대 힙, 최소 힙을 통한 풀이

import java.util.PriorityQueue;
import java.util.Scanner;

class Main {
    public static class Node {
        private int number;
        private int index;
        
        public Node(int number, int index) {
            this.number = number;
            this.index = index;
        }

        public int getNum() {
            return this.number;
        }

        public int getIndex() {
            return this.index;
        }

    }

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        for (int i = 0; i < n; i++) {
            int t = sc.nextInt();
            boolean[] visited = new boolean[t];
            PriorityQueue<Node> maxPq = new PriorityQueue<>((o1, o2) -> o2.getNum() >= o1.getNum() ? 1 : -1);
            PriorityQueue<Node> minPq = new PriorityQueue<>((o1, o2) -> o1.getNum() >= o2.getNum() ? 1 : -1);
            
            // 백준 7662 문제 풀이시 
            // (o1, o2) -> o2 - o1 과 같이 람다식을 사용하여 값을 비교하게 되면 틀렸다고 나온다.
            // 테스트 케이스에서 int의 최대값과 최소값을 비교하는 케이스가 있다고 추측된다.
            // 따라서 o2.compareTo(o1) 이나 위 방법을 사용하여 비교하여야 한다.

            for (int j = 0; j < t; j++) {
                String key = sc.next();
                int num = sc.nextInt();
                if (key.equals("I")) {
                    maxPq.add(new Node(num, j));
                    minPq.add(new Node(num, j));
                    visited[j] = true;
                } else {
                    if (num == - 1) {
                        while (!minPq.isEmpty() && !visited[minPq.peek().getIndex()]) {
                            minPq.poll();
                        } 
                        if (!minPq.isEmpty()) {
                            visited[minPq.peek().getIndex()] = false;
                            minPq.poll();
                        }
                    } else {
                        while (!maxPq.isEmpty() && !visited[maxPq.peek().getIndex()]) {
                            maxPq.poll();
                        } 
                        if (!maxPq.isEmpty()) {
                            visited[maxPq.peek().getIndex()] = false;
                            maxPq.poll();
                        }
                    }
                }
            }

            while (!minPq.isEmpty() && !visited[minPq.peek().getIndex()]) {
                minPq.poll();
            }

            while (!maxPq.isEmpty() && !visited[maxPq.peek().getIndex()]) {
                maxPq.poll();
            }

            if (!minPq.isEmpty() && !maxPq.isEmpty()) {
                System.out.println(maxPq.peek().getNum() + " " + minPq.peek().getNum());
            } else {
                System.out.println("EMPTY");
            }
        }
    }
}

위에서 제일 시간을 많이 소모했던 부분은 람다식을 통해 우선순위 큐 내부를 정리할 때였다. (o1, o2) -> o2 - o1과 같이 정렬하는 방법을 평소에 다중 정렬이 필요할때나 위 경우처럼 클래스를 직접 구현해서 문제를 풀어야 할때 애용했었는데 int 최대값과 최소값을 비교하면서 오버플로우가 난다고 한다. 따라서 compareTo 나 위와 같이 구현하여야 정답을 맞출 수 있다.

 

Java - TreeMap을 통한 풀이

import java.util.TreeMap;
import java.util.Scanner;

class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        for (int p = 0; p < n; p++) {
            int t = sc.nextInt();
            TreeMap<Integer, Integer> map = new TreeMap<>();
            for (int j = 0; j < t; j++) {
                String key = sc.next();
                int num = sc.nextInt();
                if (key.equals("I")) {
                    if (!map.containsKey(num)){
                        map.put(num, 1);
                    } else {
                        map.put(num, map.get(num) + 1);
                    }
                } else {
                    if (map.isEmpty()) continue;
                    if (num == -1) {
                        int min = map.firstKey();
                        if (map.get(min) > 1) {
                            map.put(min, map.get(min) - 1);
                        } else if (map.get(min) == 1) {
                            map.pollFirstEntry();
                        }
                    } else {
                        int max = map.lastKey();
                        if (map.get(max) > 1) {
                            map.put(max, map.get(max) - 1);
                        } else if (map.get(max) == 1){
                            map.pollLastEntry();
                        }
                    }
                }
            }
            if (!map.isEmpty()) {
                System.out.println(map.lastKey() + " " + map.firstKey());
            } else {
                System.out.println("EMPTY");
            }
        }
    }
}

TreeMap은 Map과 같이 <Key, Value>로 구성되어 있는 자료구조인데 균형 이진 트리로 부모 노드를 기준으로 작으면 왼쪽 크면 오른쪽에 배치되어 있다. 따라서 위에서는 Key에 숫자를 넣고 Value에 중복되는 수의 개수를 참조할 수 있도록 했다. 따라서 삽입 할때마다 하나씩 더해주어 전체 수의 종류와 개수를 넣어놓고 개수가 0이 되는 순간 Key:Value 를 제거하는 방식으로 정답을 구한다.

 

Python - 최대 힙, 최소 힙을 통한 풀이

from heapq import heappop, heappush
import sys
input = sys.stdin.readline

t = int(input())

for _ in range(t):
    k = int(input())
    max_heap = []
    min_heap = []
    visited = [False] * (1000001)
    for i in range(k):
        char, n = map(str, input().split())
        n = int(n)
        if char == "I":
            heappush(min_heap, (n, i))
            heappush(max_heap, (-n, i))
            visited[i] = True
        else:
            if n == -1:
                while min_heap and not visited[min_heap[0][1]]:
                    heappop(min_heap)
                if min_heap:
                    visited[min_heap[0][1]] = False
                    heappop(min_heap)
            else:
                while max_heap and not visited[max_heap[0][1]]:
                    heappop(max_heap)
                if max_heap:
                    visited[max_heap[0][1]] = False
                    heappop(max_heap)
    
    while min_heap and not visited[min_heap[0][1]]: 
        heappop(min_heap)
    while max_heap and not visited[max_heap[0][1]]:
        heappop(max_heap)

    if min_heap and max_heap:
        print(-max_heap[0][0], min_heap[0][0])
    else:
        print("EMPTY")

기본적인 내용은 위의 Java 최대 힙, 최소 힙을 통한 풀이방법과 동일하다. 하지만 Python의 heapq에서는 Comparator를 바꿔서 정렬 순서를 최대힙에 맞게 바꿀 수 없으므로 원래의 숫자에 -1을 곱하여 최대 힙으로 바뀌게끔 하였다. 최대 힙으로 바꾼 수는 연산이 끝나고 마지막에 다시 -1을 곱해 원래의 수가 되게끔 만들어줬다.

 

정리

Python에서 최대 힙을 구현하려면  -1을 곱해주기만 하면 된다는 걸 잊지말자. 또한 튜플은 값을 변경할 수 없지만 위와 같이 한번 넣어놓고 다시는 바뀌지 않을 값이라면 리스트 생성보다 빨리 동작하기 때문에 튜플을 통하여 값을 지정하자. Java에서 람다식을 통한 (o1, o2) -> o1 - o2 비교 시 오버플로우가 일어날 수 있다. 이 점을 생각해서 다른 구현 방법을 익혀두자.