Untitled
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