Post

[Python] 이중 우선순위 큐

💡 우선순위 큐란?

우선순위 큐(Priority Queue, 줄여서 pq)는 가장 우선순위가 높은(혹은 낮은) 요소를 빠르게 꺼낼 수 있는 자료구조다. 힙으로 구현한다.
파이썬에서는 heapq 모듈을 통해 최소 힙(min-heap)을 기본으로 사용할 수 있다.

  • heapq.heappush(): 삽입 연산 → O(log N)
  • heapq.heappop(): 삭제 연산 → O(log N)



문제

백준 7662. 이중 우선순위 큐

이 문제에서는 최댓값과 최솟값을 빠르게 꺼낼 수 있어야 하기 때문에 우선순위 큐를 사용해야 한다.


🤔 왜 힙을 두 개 사용해야 할까? (min-heap, max-heap)

Python의 heapq는 최소 힙만 지원한다. 하지만 이 문제는 최댓값과 최솟값 모두를 빠르게 삭제해야 하기 때문에, 최대 힙도 따로 필요하다.

최소 힙에 음수 부호를 붙여서 최대 힙처럼 동작하게 할 수 있다.

1
heapq.heappush(max_heap, -1 * num)

이렇게 값을 음수로 바꿔서 넣은 다음 나중에 pop할 때는 다시 음수 부호를 붙여서 원래 값으로 되돌리면 된다. 큰 수일수록 음수 기준으로는 작게 취급되므로, heapq에서 가장 큰 수를 빼낸 것처럼(최대 힙인 것처럼) 보이게 된다.

1
2
heapq.heappush(min_heap, num)
heapq.heappush(max_heap, -1 * num)

즉, 이렇게 최소 힙 하나, 최대 힙 하나를 따로 운용해서
각각의 역할(최솟값, 최댓값 추출)을 처리하는 것이 이 문제의 핵심이다.


🧨 한계: 특정 값 삭제가 불가능함

이 문제에서는 최대/최소 값을 삭제하는 연산도 구현해야 한다.
최댓값 삭제 시 최대 힙에서 pop하고, 최솟값 삭제 시 최소 힙에서 pop하면 되는데, 이렇게 하면 최대 힙이랑 최소 힙이 서로 달라져버리기 떄문에 그걸 매번 맞춰줘야 했다.

처음에는, ‘최대 힙에서 pop한 숫자를 최소 힙에서도 지워주면 되지 않나?’라고 생각했었다. 근데 제약이 있었다. Python의 heapq는 특정 값을 직접 삭제하는 기능을 지원하지 않는다.

왜 그럴까?
힙 자료구조는 루트 노드(최솟값 or 최댓값) 기준으로 빠르게 삽입하고 꺼내는 데에 최적화되어 있다. 그러니 임의의 값을 찾아서 삭제하는 작업은 힙 구조의 핵심 철학과 어긋나는 거 아닐까?(거창…)

1
2
3
heap.remove(x)   # O(N) → 비효율적
heapq.heapify(heap)  # O(N) → 또 비효율적
> 시간 초과!

달리 떠오르는 방법이 없어 이렇게 작성했더니 시간초과가 났다. 이런 방식은 힙을 사용하는 이유 자체를 무너뜨리는 비효율적인 방식이다.

정답은 힙에서 직접 값을 지우지 않고, ‘이미 삭제된 값은 건너뛰는 방식’으로 우회하는 것이었다.


✅ 해결 방법: valid 딕셔너리로 동기화

valid 딕셔너리를 만들고, 특정 숫자가 현재 힙에 몇 개 남아있는지 기록한다. 그리고 요소를 삽입하거나 삭제할 때마다 이런 작업을 해준다.

  • 요소 삽입 시:
    1
    2
    3
    
    heapq.heappush(min_heap, num)
    heapq.heappush(max_heap, -1*num)
    valid[num] = valid.get(num, 0) + 1  # 남은 개수에 1을 더해줌
    
  • 요소 삭제 시:
    1
    2
    3
    4
    5
    
     while min_heap:
      d = heapq.heappop(min_heap)
      if valid[d] > 0:  # 1개 이상 남아있다면
          valid[d] -= 1  # 남은 개수에 1을 빼줌
          break
    

valid에 적힌 값을 확인해서, 유효한 숫자가 나올 때까지(valid값이 0이라면 이미 없는 숫자이므로 다음 숫자를 봐야한다) 계속 pop한다. 유효한 값이라면 valid 값에서 1을 뺀 다음 break 해준다.



🧾 전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import sys
import heapq

T = int(input())
for _ in range(T):
    k = int(input())
    min_heap = []
    max_heap = []
    valid = {}
    for _ in range(k):
        q, num = sys.stdin.readline().split()
        num = int(num)
        if q == 'I':  # 삽입
            heapq.heappush(min_heap, num)
            heapq.heappush(max_heap, -1*num)
            valid[num] = valid.get(num, 0) + 1
        elif q == 'D':  # 삭제
            if num == 1:  # 최대값 삭제
                while max_heap:
                    d = -1*heapq.heappop(max_heap)
                    if valid[d] > 0:
                        valid[d] -= 1
                        break
            else:  # 최소값 삭제
                while min_heap:
                    d = heapq.heappop(min_heap)
                    if valid[d] > 0:
                        valid[d] -= 1
                        break

    # valid에 0인 숫자가 힙에 남아있다면 제거
    while max_heap and valid[-1*max_heap[0]] == 0:
        heapq.heappop(max_heap)
    while min_heap and valid[min_heap[0]] == 0:
        heapq.heappop(min_heap)

    if min_heap and max_heap:
        print(-1*max_heap[0], min_heap[0])
    else:
        print('EMPTY')
This post is licensed under CC BY 4.0 by the author.

Trending Tags