이중 우선순위 큐

다른 블로그에서 이중 우선순위 큐 C++ 포스팅을 보고 파이썬으로 풀어봤는데 생각보다 굉장히 어려웠다.

I 명령이 들어오면 원소를 추가하고, D 명령이 들어오면 ‘1’일 때는 최대값을 삭제하고, ‘-1’일 때는 최소값을 삭제한다. 힙을 사용해 원소들을 정렬하되 최소값과 최대값 모두를 추적해야 해서 단순히 힙을 써서는 풀지 못한다. 힙 자료구조는 왼쪽 끝의 최소값은 보장하지만, 최대값은 보장하지 않기 때문이다.

import heapq
a=[1,3,4,5,3,4]
heapq.heapify(a)
print(a)

[1,3,4,5,3,4] 리스트를 heap으로 변환시켜주는 heapify를 실행 후 출력해보면 아래와 같다. 첫 번째로 위치한 값은 최소값이 맞지만, 마지막으로 위치한 값은 이 리스트의 최대값이 아님을 확인할 수 있다.

[1, 3, 4, 5, 3, 4]

따라서, 이 문제에서는 입력에 대해 최대힙과 최소힙을 모두 사용하되 두 힙을 서로 연동시켜주기 위한 ‘key’도 함께 사용해야 한다.

  • 먼저 최대힙, 최소힙을 각각 하나씩 만들어 준다.
  • 최대힙은 최소힙에서 원소에 ‘-‘를 붙이면 만들 수 있다.
  • key 로 사용할 리스트를 선언하고 값들을 true로 초기화 해준다.
  • 원소를 삭제할 때 최소값은 최소힙에서 최대값은 최대힙에서 삭제하고, 삭제된 값의 key를 false로 바꿔준다.
  • 최소힙에서 원소를 삭제할 때 삭제한 원소가 이미 최대힙에서 삭제된 값일 수도 있음으로, 항상 key 값을 확인하며 key가 true인 원소가 나올 때 까지 원소를 없애준다. 최대힙에서 원소를 삭제할 때도 마찬가지.
  • 연산을 모두 마치고 최대값, 최소값을 출력하기 전에 아직 key 값이 false인 다른 힙에서 이미 삭제된 원소가 남아있을 수 있음으로 최대값 혹은 최소값이 모두 true가 되도록 최대힙, 최소힙에서 원소를 삭제해준다.


CODE

import sys
import heapq

for _ in range(int(input())):
    k=int(input())
    ids=list(sys.stdin.readline().split() for _ in range(k))
    for i in range(len(ids)):
        ids[i][1]=int(ids[i][1])

    min_hq=[]
    max_hq=[]
    key=[False]*1000001

    for index,id in enumerate(ids):
        if id[0]=='I':
            heapq.heappush( min_hq, (id[1],index))
            heapq.heappush( max_hq, (-id[1],index))
            key[index]=True

#최소값의 key가 True가 나오거나 리스트가 빌 때까지 원소를 뽑아낸다.
#이렇게 동기화가 완료되었으면 최소값 하나를 뽑아내고, 그 원소의 key를 false로 변경
        elif id[1] == -1:
            while min_hq and not key[min_hq[0][1]]:
                heapq.heappop(min_hq)
            if min_hq:
                key[heapq.heappop(min_hq)[1]]=False

#최대값의 key가 True가 나오거나 리스트가 빌 때까지 원소를 뽑아낸다.
#이렇게 동기화가 완료되었으면 최대값 하나를 뽑아내고, 그 원소의 key를 false로 변경
        elif id[1] == 1:
            while max_hq and not key[max_hq[0][1]]:
                heapq.heappop(max_hq)
            if max_hq:
                key[heapq.heappop(max_hq)[1]]=False

#최대값 혹은 최소값이 이미 다른 힙에서 삭제된 원소일 수도 있기 때문에
#key가 true인 원소가 나올 때 까지 최대값, 최소값을 삭제한다.
    while min_hq:
        if not key[min_hq[0][1]]:
            heapq.heappop(min_hq)
        else: break

    while max_hq:
        if not key[max_hq[0][1]]:
            heapq.heappop(max_hq)
        else: break

    if not min_hq or not max_hq: print('EMPTY')
    else:
        print(-heapq.heappop(max_hq)[0], heapq.heappop(min_hq)[0])


[다시 풀어보기]