[프로그래머스 Lv.3] 이중우선순위큐(Java) - TreeMap과 두 개의 힙

📌 핵심 요약

• 힙 하나로는 안 된다 — PriorityQueue는 루트 한쪽 끝만 정렬을 보장한다. 최소 힙에서 최댓값의 위치는 알 수 없다.
• 풀이 1: TreeMap — 값을 키, 개수를 밸류로 저장한다. firstKey와 lastKey로 양쪽 끝을 각각 O(log n)에 얻는다.
• 풀이 2: 두 개의 힙 — 최소 힙과 최대 힙에 같은 값을 넣고, 한쪽에서 꺼낸 값을 다른 쪽에서 remove로 지운다.
• 차이는 삭제 비용 — PriorityQueue.remove(Object)는 선형 탐색이라 O(n)이다. TreeMap 쪽이 최악의 경우에 안전하다.

 

목차

1 문제

프로그래머스 #42628

이중우선순위큐

operations.length ≤ 1,000,000

난이도: Lv.3

최댓값과 최솟값을 모두 꺼낼 수 있는 큐를 구현하는 문제다. 명령은 "명령어 데이터" 형식의 문자열로 주어진다.

명령 종류

1. I 숫자 — 해당 숫자를 큐에 삽입한다
2. D 1 — 큐에서 최댓값을 삭제한다
3. D -1 — 큐에서 최솟값을 삭제한다

놓치면 안 되는 조건

1. 최댓값이나 최솟값이 둘 이상이면 하나만 삭제한다
2. 빈 큐에 삭제 명령이 오면 그 연산은 무시한다
3. 모든 연산 후 큐가 비어 있으면 [0, 0], 아니면 [최댓값, 최솟값]을 반환한다

2 힙 하나로는 왜 안 되는가

처음에는 PriorityQueue 하나로 시작했다. D -1은 poll() 한 번이면 끝나니 절반은 해결된 셈이었다. 그런데 D 1에서 막혔다.

   

Java — 처음 작성한 뼈대

if (num == 1) {
    // 최댓값을 어떻게 꺼내지?
} else {
    pq.poll();
}

힙은 부모가 자식보다 작다는 조건만 유지하는 구조다. 형제끼리는 아무 순서도 없고, 서로 다른 가지에 있는 노드 사이에도 순서가 없다. 최소 힙 [1, 3, 2, 9, 5, 8, 4]를 트리로 그려 보면 이렇다.

깊이	노드	보장되는 것
0 (루트)	1	전체 최솟값
1	3, 2	각자 부모인 1보다 크다는 것만
2 (리프)	9, 5, 8, 4	아무 순서도 없음

최댓값 9는 리프 어딘가에 있다는 것만 알 수 있다. 정확한 위치를 알아내려면 결국 배열 전체를 훑어야 하고, D 1이 나올 때마다 이 작업을 반복하면 최악의 경우 연산량이 operations.length의 제곱에 비례한다. 제한이 100만이므로 감당할 수 없다.

ℹ️

필요한 조건

양쪽 끝을 모두 로그 시간에 다룰 수 있는 구조가 필요하다. 여기서 방향이 두 갈래로 갈린다. 하나는 애초에 양쪽 끝을 아는 구조를 쓰는 것(TreeMap), 다른 하나는 한쪽 끝만 아는 구조를 두 개 두는 것(두 개의 힙)이다.

3 풀이 1 — TreeMap

구조 설계

TreeMap은 키를 정렬된 상태로 유지하는 맵이다. firstKey()는 가장 작은 키를, lastKey()는 가장 큰 키를 각각 O(log n)에 반환한다. 하나의 구조가 양쪽 끝을 모두 안다.

다만 맵은 같은 키를 두 번 담을 수 없다. I 5가 두 번 오면 원소가 두 개여야 하는데 키는 하나뿐이다. 그래서 값을 키로, 그 값이 몇 개 들어 있는지를 밸류로 저장한다. 삭제할 때는 키를 지우는 것이 아니라 개수를 줄이고, 개수가 0이 되는 순간에만 키를 제거한다.

동작 추적

["I 7", "I 5", "I 5", "D 1", "D -1"]로 상태 변화를 따라갔다.

명령	하는 일	맵 상태	원소 수
시작	—	{}	0
I 7	7의 개수 +1	{7:1}	1
I 5	5의 개수 +1	{5:1, 7:1}	2
I 5	5의 개수 +1	{5:2, 7:1}	3
D 1	lastKey=7의 개수 -1 → 0이므로 키 제거	{5:2}	2
D -1	firstKey=5의 개수 -1 → 1이 남아 키 유지	{5:1}	1

마지막에 [lastKey, firstKey]는 [5, 5]가 된다. 원소가 하나뿐일 때 최댓값과 최솟값이 같아지는 경우가 따로 처리하지 않아도 맞아떨어진다.

코드

   

Java — 풀이 1 (TreeMap)

import java.util.*;

class Solution {
    public int[] solution(String[] operations) {
        // 키: 값 자체 / 밸류: 그 값이 큐에 들어 있는 개수
        TreeMap<Integer, Integer> counts = new TreeMap<>();

        for (String operation : operations) {
            String[] parts = operation.split(" ");
            String order = parts[0];
            int num = Integer.parseInt(parts[1]);

            if (order.equals("I")) {
                // 처음 등장하는 값이면 기본값 0에서 시작
                counts.put(num, counts.getOrDefault(num, 0) + 1);
            } else {
                // 빈 큐에 대한 삭제 명령은 무시
                if (counts.isEmpty()) {
                    continue;
                }
                int target = (num == 1) ? counts.lastKey() : counts.firstKey();
                int count = counts.get(target);

                if (count > 1) {
                    counts.put(target, count - 1);
                } else {
                    counts.remove(target);
                }
            }
        }

        if (counts.isEmpty()) {
            return new int[]{0, 0};
        }
        return new int[]{counts.lastKey(), counts.firstKey()};
    }
}

💡

getOrDefault를 한쪽에만 쓴 이유

I 분기의 num은 맵에 없을 수 있으므로 기본값 0이 제 역할을 한다. 반면 D 분기의 target은 방금 lastKey()나 firstKey()로 꺼낸 값이라 반드시 존재한다. 여기에도 getOrDefault를 쓰면 읽는 사람이 "키가 없을 수도 있나?" 하고 멈추게 된다. 같은 메서드라도 필요한 자리와 아닌 자리를 구분하면 코드가 스스로 사실을 설명한다.

4 풀이 2 — 최소 힙 + 최대 힙

구조 설계

힙 하나가 한쪽 끝만 안다면, 힙을 두 개 두면 양쪽 끝을 다 알 수 있다. 삽입할 때 같은 값을 두 힙에 모두 넣는다. 그러면 두 힙은 항상 동일한 원소 집합을 갖고, 보는 방향만 다르다.

문제는 삭제다. maxHeap.poll()로 최댓값을 꺼내면 minHeap에는 그 값이 그대로 남는다. 두 힙의 내용이 어긋나므로 다른 쪽에서도 같은 값을 지워 줘야 한다. 이때 remove(Object)를 쓴다.

동작 추적

같은 입력 ["I 7", "I 5", "I 5", "D 1", "D -1"]로 두 힙의 상태를 따라갔다. 힙 내부 배열 순서는 구현에 따라 달라지므로 여기서는 담긴 값의 집합으로 표기했다.

명령	하는 일	minHeap	maxHeap
I 7	양쪽에 7 삽입	{7}	{7}
I 5	양쪽에 5 삽입	{5, 7}	{5, 7}
I 5	양쪽에 5 삽입	{5, 5, 7}	{5, 5, 7}
D 1	maxHeap에서 7을 poll → minHeap에서 7을 remove	{5, 5}	{5, 5}
D -1	minHeap에서 5를 poll → maxHeap에서 5를 remove	{5}	{5}

마지막 D -1에서 중복 값 5가 두 개 있었지만 결과는 정확하다. remove(Object)는 같은 값이 여러 개여도 하나만 제거하기 때문이다. 문제 조건인 "최댓값이 둘 이상이면 하나만 삭제한다"와 정확히 맞물린다.

코드

   

Java — 풀이 2 (두 개의 힙)

import java.util.*;

class Solution {
    public int[] solution(String[] operations) {
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

        for (String op : operations) {
            String[] tokens = op.split(" ");
            String command = tokens[0];
            int value = Integer.parseInt(tokens[1]);

            if (command.equals("I")) {
                // 두 힙이 같은 원소 집합을 갖도록 양쪽에 삽입
                minHeap.add(value);
                maxHeap.add(value);
            } else if (command.equals("D")) {
                // 두 힙의 크기는 항상 같으므로 한쪽만 검사하면 된다
                if (minHeap.isEmpty()) continue;

                if (value == 1) {
                    int max = maxHeap.poll();
                    minHeap.remove(max);
                } else {
                    int min = minHeap.poll();
                    maxHeap.remove(min);
                }
            }
        }

        if (minHeap.isEmpty()) {
            return new int[]{0, 0};
        }
        return new int[]{maxHeap.peek(), minHeap.peek()};
    }
}

⚠️

remove(Object)와 remove(int)의 구분

minHeap.remove(max)에서 max는 int지만, PriorityQueue에는 인덱스로 지우는 remove(int)가 없다. 따라서 Collection.remove(Object)가 호출되고 max는 Integer로 오토박싱된 뒤 equals로 비교된다. 값 기준 삭제이므로 의도대로 동작한다.
다만 같은 코드를 ArrayList에 쓰면 인덱스 삭제인 remove(int)가 선택되어 전혀 다른 원소가 지워진다. 자료구조를 바꿀 때 이 줄은 반드시 다시 봐야 한다.

5 두 풀이 비교

두 풀이 모두 정답을 낸다. 차이는 삭제 한 번의 비용에서 갈린다.

항목	풀이 1 (TreeMap)	풀이 2 (두 개의 힙)
삽입	O(log n)	O(log n) × 2회
최댓값 조회	lastKey() — O(log n)	peek() — O(1)
삭제	O(log n)	poll은 O(log n), remove는 O(n)
최악 전체	O(n log n)	O(n²)
중복 값 처리	개수 밸류로 관리	원소를 그대로 여러 개 보관
동기화 필요	없음 (구조 하나)	있음 (두 힙을 맞춰야 함)
메모리	서로 다른 값의 개수만큼	전체 원소 × 2

remove가 O(n)인 이유

2번 섹션에서 본 것과 같은 이유다. 힙에서 특정 값이 어디 있는지는 루트가 아닌 이상 알 수 없다. PriorityQueue.remove(Object)는 내부 배열을 앞에서부터 선형으로 훑어 일치하는 원소를 찾은 뒤, 그 자리를 마지막 원소로 채우고 힙 조건을 복구한다. 찾는 과정이 O(n)이라 전체가 O(n)이 된다.

operations.length가 100만이고 그중 절반이 I, 절반이 D라면 힙에 최대 50만 개가 쌓인 상태에서 50만 번의 선형 탐색이 일어난다. 이론상 10¹¹ 수준이다. 

🚨

통과 여부와 별개로

두 힙 풀이는 발상이 직관적이라 떠올리기 쉽지만, 최악의 경우를 계산해 보면 제한을 감당하지 못하는 구조다. 제출이 통과하더라도 "왜 통과했는가"가 채점 데이터에 달려 있다면 근거로 삼기 어렵다. 제한사항을 먼저 계산하고 자료구조를 고르는 순서를 지키면 이 판단이 코드를 쓰기 전에 나온다.

두 힙을 살리려면

두 힙 구조를 유지하면서 remove를 없애는 방법도 있다. 지연 삭제(lazy deletion)다. 삭제 시 다른 쪽 힙은 건드리지 않고 "이 원소는 무효"라는 표시만 남긴 뒤, 나중에 poll할 때 무효 표시된 원소를 걷어내며 진행한다.

다만 이때는 값이 아니라 삽입 순서 인덱스를 함께 담아야 한다. 같은 값이 여러 개일 때 어느 것이 무효인지 구분해야 하기 때문이다. 힙에 넣을 원소가 int에서 (값, 인덱스) 쌍으로 바뀌고, 별도의 boolean[] 배열을 동기화해야 한다. 성능은 O(n log n)으로 개선되지만 관리할 상태가 셋으로 늘어난다.

TreeMap은 이 동기화 문제 자체가 생기지 않는다. 구조가 하나라 어긋날 대상이 없다. 자료구조가 실수 가능성을 줄여 주는 쪽을 고른다는 기준으로 보면 풀이 1이 낫다고 판단했다.

6 실수했던 부분

힙 하나로 최댓값을 다루려 했다

원인. PriorityQueue를 "정렬된 큐"로 생각했다. 실제로는 루트만 정렬을 보장하는 부분 순서 구조이고, 나머지는 부모-자식 관계만 유지된다. 반대쪽 끝을 상수 시간에 볼 방법이 없다.

해결. D 1 분기를 비워 둔 채로 넘어가지 않고, 그 자리에서 "이 구조가 최댓값의 위치를 아는가"를 먼저 확인했다. 답이 아니라는 것이 확인되자 자료구조 선택으로 문제가 옮겨 갔다.

잘못된 내용이 있다면 지적부탁드립니다. 방문해주셔서 감사합니다.