[7662] 이중 우선순위 큐 - Python

/* 문제 */
이중 우선순위 큐(dual priority queue)는 전형적인 우선순위 큐처럼 데이터를 삽입, 삭제할 수 있는 자료 구조이다. 전형적인 큐와의 차이점은 데이터를 삭제할 때 연산(operation) 명령에 따라 우선순위가 가장 높은 데이터 또는 가장 낮은 데이터 중 하나를 삭제하는 점이다. 이중 우선순위 큐를 위해선 두 가지 연산이 사용되는데, 하나는 데이터를 삽입하는 연산이고 다른 하나는 데이터를 삭제하는 연산이다. 데이터를 삭제하는 연산은 또 두 가지로 구분되는데 하나는 우선순위가 가장 높은 것을 삭제하기 위한 것이고 다른 하나는 우선순위가 가장 낮은 것을 삭제하기 위한 것이다. 
정수만 저장하는 이중 우선순위 큐 Q가 있다고 가정하자. Q에 저장된 각 정수의 값 자체를 우선순위라고 간주하자. 
Q에 적용될 일련의 연산이 주어질 때 이를 처리한 후 최종적으로 Q에 저장된 데이터 중 최댓값과 최솟값을 출력하는 프로그램을 작성하라.

최댓값과 최솟값을 삭제하는 연산을 해야 하므로 최대 힙과 최소 힙 두 개의 힙을 사용했다. 힙을 두 개 사용하여 각 연산마다 값을 업데이트 한다는 아이디어까지는 어렵지 않게 떠올릴 수 있었는데, 문제는 이 두 힙을 동기화하는 것이었다.

 

처음에는 리스트의 remove()를 사용하여 최댓값을 삭제할 때는 최소 힙에서 그 값을 같이 삭제해주고, 최솟값을 삭제할 때는 최대 힙에서 그 값을 같이 삭제해줬다.

# 시간 초과 코드
import sys
input = sys.stdin.readline
import heapq

t = int(input())
for _ in range(t):
    minQ = []
    maxQ = []
    k = int(input())
    for i in range(k):
        com, n = input().split()
        n = int(n)
        if com == 'I':
            heapq.heappush(minQ, n)
            heapq.heappush(maxQ, -n)
        else:
            if len(minQ) > 0:
                if n == 1:
                    d = -heapq.heappop(maxQ)
                    minQ.remove(d)
                    heapq.heapify(minQ)
                else:
                    d = heapq.heappop(minQ)
                    maxQ.remove(-d)
                    heapq.heapify(maxQ)
                    
    if len(minQ) == 0:
        print('EMPTY')
    else:
        print(-heapq.heappop(maxQ), heapq.heappop(minQ))

그러나 이 코드는 시간 초과가 떴다. 매번 리스트에서 값을 삭제하고 heapify를 해주는 과정에서 시간이 오래 걸린 것 같다.

 

그래서 생각한 방법은 별도의 리스트를 생성해 특정 노드가 삭제되었는지 아닌지를 저장하는 것이었다.

# 정답 코드
import sys
input = sys.stdin.readline
import heapq

t = int(input())
for _ in range(t):
    k = int(input())
    minQ, maxQ = [], []
    # 각 노드의 id에 대해 삭제되었는지 아닌지를 저장
    # 처음에는 아무 값도 없으므로 모두 삭제된 것으로 간주
    deleted = [True] * k  
    for i in range(k):
        com, n = input().split()
        n = int(n)
        if com == 'I':
            heapq.heappush(minQ, (n, i))
            heapq.heappush(maxQ, (-n, i))
            deleted[i] = False
        else:
            if n == 1:
                # 삭제되지 않은 값 찾기
                # 삭제된 값은 힙에서 제거
                while maxQ and deleted[maxQ[0][1]]:
                    heapq.heappop(maxQ)
                if maxQ:
                    deleted[maxQ[0][1]] = True
                    heapq.heappop(maxQ)
            else:
                while minQ and deleted[minQ[0][1]]:
                    heapq.heappop(minQ)
                if minQ:
                    deleted[minQ[0][1]] = True
                    heapq.heappop(minQ)

    # 연산이 끝난 후 삭제된 값들 제거
    while minQ and deleted[minQ[0][1]]:
        heapq.heappop(minQ)
    while maxQ and deleted[maxQ[0][1]]:
        heapq.heappop(maxQ)
    
    if minQ and maxQ:
        print(-maxQ[0][0], minQ[0][0])
    else:
        print('EMPTY')

deleted 배열은 각 노드의 아이디에 따라 그 노드가 삭제되었는지 아닌지를 저장한다. 노드의 id는 반복문을 몇 번째 도는지, 즉 i 변수의 값으로 지정했다. 힙에 값을 넣을 때도 이 id를 튜플 형태로 함께 넣어주었는데, 이렇게 하면 모든 노드에 유일한 id를 지정할 수 있다. id 값을 지정함으로써 이번 연산에 사용된 노드가 삭제 대상인지 삽입 대상인지를 구분할 수 있는 것이다. 삭제 대상이었다면 그 노드는 다음 반복문에서 삭제되고, 삽입 대상이었다면 다음 반복문에서 최댓값/최솟값을 구할 때 포함된다.

모든 연산이 끝난 후 아직 힙에서 제거되지 않은 삭제 대상 노드들을 모두 지워주면 최종적으로 두 개의 힙이 동기화된다.