Untitled

mail@pastecode.io avatar
unknown
python
8 months ago
2.1 kB
3
Indexable
Never
import heapq


class BiHeap:
    def __init__(self, capacity):
        self.min_heap = []
        self.max_heap = []
        self.cnt = 0
        self.capacity = capacity
        self.deleted = set()
        self.timestamp = 0

    def get_min(self):
        self.clean_min()

        return self.min_heap[0] if self.min_heap else None

    def clean_min(self):
        while self.min_heap and self.min_heap[0] in self.deleted:
            self.deleted.discard(self.min_heap[0])
            heapq.heappop(self.min_heap)

    def clean_max(self):
        while self.max_heap and (-self.max_heap[0][0], -self.max_heap[0][1], -self.max_heap[0][2]) in self.deleted:
            self.deleted.discard(
                (-self.max_heap[0][0], -self.max_heap[0][1], -self.max_heap[0][2]))
            heapq.heappop(self.max_heap)

    def get_max(self):
        if self.max_heap is None:
            return None

        a, b, c = self.max_heap[0]
        return -a, -b, -c

    def pop_min(self):
        a, b, c = heapq.heappop(self.min_heap)
        self.deleted.add((a, b, c))
        return a, b, c

    def pop_max(self):
        a, b, c = heapq.heappop(self.max_heap)
        self.deleted.add((-a, -b, -c))
        return -a, -b, -c

    def insert(self, a, b, c):
        heapq.heappush(self.min_heap, (a, b, c))
        heapq.heappush(self.max_heap, (-a, -b, -c))
        self.cnt += 1

        if self.cnt > self.capacity:
            self.pop_min()
            self.cnt -= 1

    def enqueue(self, val, priority):
        self.insert(priority, -self.timestamp, val)
        self.timestamp += 1

    def dequeue(self):
        neg_priority, timestamp, val = self.pop_max()
        priority = -neg_priority
        self.cnt -= 1

        return val, priority


biheap = BiHeap(3)
biheap.enqueue(5, 3)
biheap.enqueue(6, 5)
biheap.enqueue(1, -1)
biheap.dequeue()
biheap.enqueue(-2, 0)
biheap.dequeue()
biheap.enqueue(3, 1)
biheap.enqueue(4, 2)


while biheap.get_min() is not None:
    print(biheap.dequeue())
Leave a Comment