Priority Queue

Heapq만을 사용한 Priority Queue의 Time Complexity는 아래와 같다.

Find max/min - O(1)

Insert - O(logN)

Remove - O(N)

Remove Operation에 Linear Time이 필요하다는 것이 아쉽다.

 

이를 해결하기 위한 Custom Priority Queue를 만들었다.

이는 두 개의 Priority Queue를 사용하는 방식이다.

하나의 Queue는 원래 목적인 max/min heap로 사용하고, 나머지 하나의 queue를 쓰레기통으로 사용한다.

Time Complexity는 아래와 같다.

Find max/min - O(1)

Insert - O(logN)

Remove - O(logN)

 

다른 연산은 그다지 어렵지 않으니 Remove 연산만 간단히 살펴보자.

삭제하려는 원소에 따라 크게 두 경우가 존재한다.

1) 삭제하려는 원소가 max heap의 top에 있다.

A - 이러면 max heap의 top 원소를 날리고,

B - 쓰레기통의 top과 max heap의 top이 같은 경우 원소를 서로 삭제한다.

쓰레기통의 top과 max heap의 top이 달라질 때까지 B를 반복한다.

2) 삭제하려는 원소가 max heap의 top에 없다.

현재 삭제 불가능하므로 쓰레기통에 삽입해둔다.

import heapq

class PriorityQueue:
    def __init__(self):
        self.pq = []
        self.removals = []

    def insert(self, elem):
        heapq.heappush(self.pq, -elem)

    def gettop(self):
        return -self.pq[0]

    def remove(self, elem):
        if elem == -self.pq[0]:
            heapq.heappop(self.pq)
            while self.removals and self.pq[0] == self.removals[0]:
                heapq.heappop(self.pq)
                heapq.heappop(self.removals)
        else:
            heapq.heappush(self.removals, -elem)

    def length(self):
        return len(self.pq)