프로그래머스: 이중우선순위큐 [JAVA]

문제 링크

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

코딩테스트 연습 - 이중우선순위큐

풀이

전체적인 풀이 과정은 다음과 같다.

 

1. 최대 힙, 최소 힙 생성
2. 주어진 operations 배열들에 대해 각 연산을 적용
   • I 명령인 경우 힙 2개 모두에 값을 추가
   • D 1 명령(최대값 삭제)인 경우 최대 힙에서는 poll() 연산을 진행, 최소 힙에서는 remove() 연산을 진행
   • D -1 명령(최대값 삭제)인 경우 최소 힙에서는 poll() 연산을 진행, 최대 힙에서는 remove() 연산을 진행
   • 단, 여태까지 누적된 삽입 명령어의 개수가 삭제 명령어의 개수보다 클때에만 삭제 명령어를 진행(아니면 건너뛴다.)
3. 힙이 비었다면 (누적 삽입 명령어 개수 == 누적 삭제 명령어 개수) 정답 배열을 0, 0 처리하고, 아니면 각 힙에서 하나씩 poll()해서 값을 초기화

 해당 문제의 경우엔 고민과 다르게 해결이 매우 허무하게 끝난 문제이다. 문제의 이름처럼 주어진 숫자들에서 최대값과 최소값을 최소한의 복잡도로 뺄 수 있는 알고리즘을 작성해야 하는데, 기존의 자료구조들로는 불가능하게 여겨져 이를 어떻게 구현해야 하나를 고민했다.

 그러다 도저히 방도가 생각나지 않아 기존의 힙(자바에서는 우선순위 큐)을 통해 최소 힙, 최대 힙을 만들어 그냥 양쪽에서 번갈아 작업을 진행하는 것으로 작성했다. 물론 처음에는 한쪽 힙에서 제거한 수를 다른 쪽에선 어떻게 제거해야 하는가를 매우 고민했는데, 후에 찾아보니 우선순위큐에서 단순히 가장 최우선순위 수를 제거하는 함수(poll()) 뿐만 아니라 그냥 하나의 객체를 제거하는 함수(remove())를 제공해주는 것을 뒤늦게 알고, 이를 적용해 문제를 매우 쉽게 해결했다.

 결국 최대 힙, 최소 힙을 각각 선언한 다음, 주어진 operations 배열을 순서대로 처리를 하면 된다. I 명령(숫자 삽입 명령)의 경우 단순히 양쪽 힙에 값을 모두 추가하면 되고, D의 경우 최대냐 최소냐에 따라 해당 힙에서 poll() 명령을 통해 값을 제거한 다음, 다른 쪽의 힙에 remove() 명령을 통해 해당 숫자를 제거해주면 된다.

 이 때 중요한 것은 삽입 명령과 삭제 명령의 개수를 모두 저장하여, 비어있는 힙에서 제거 명령을 내리는 상황(삭제 명령의 개수 >= 삽입 명령의 개수)에서는 해당 명령 자체를 무시하도록 코드를 작성하면 된다.

 최종적으로는 힙이 비어있는 상황을 제외하고는, 최대 힙과 최소 힙에서 수를 하나씩 빼와 answer 배열을 초기화하면 된다.

 

class Solution {
    public int[] solution(String[] operations) {
        int[] answer = new int[2];
        int insertCount = 0, deleteCount = 0;
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder()); //최대 힙
        PriorityQueue<Integer> minHeap = new PriorityQueue<>(); //최소 힙

        for (int i = 0; i < operations.length; i++) { //각 명령
            String[] operation = operations[i].split(" ");

            String instruction = operation[0]; //명령어
            int num = Integer.parseInt(operation[1]); //숫자 (or 1/-1)

            if (instruction.equals("I")) { //삽입 명령인 경우
                ++insertCount;

                maxHeap.add(num);
                minHeap.add(num);

            } else { //삭제 명령인 경우
                if (insertCount > deleteCount) { //삽입 명령이 더 많은 상태에서만
                    ++deleteCount;

                    if (num == 1) {//최대값 제거인 경우
                        Integer deleteNum = maxHeap.poll();
                        minHeap.remove(deleteNum);
                    } else {//최소값 제거인 경우
                        Integer deleteNum = minHeap.poll();
                        maxHeap.remove(deleteNum);
                    }
                }
            }
        }

        if (insertCount == deleteCount) {
            answer[0] = 0;
            answer[1] = 0;
        } else {
            answer[0] = maxHeap.poll();
            answer[1] = minHeap.poll();
        }
        return answer;
    }
}

결과