from collections import MutableMapping class _MinEntry(object): """ Mutable entries for a Min-PQ dictionary. """ def __init__(self, dkey, pkey): self.dkey = dkey #dictionary key self.pkey = pkey #priority key def __lt__(self, other): return self.pkey < other.pkey class _MaxEntry(object): """ Mutable entries for a Max-PQ dictionary. """ def __init__(self, dkey, pkey): self.dkey = dkey self.pkey = pkey def __lt__(self, other): return self.pkey > other.pkey class PQDict(MutableMapping): def __init__(self, *args, **kwargs): self._heap = [] self._position = {} self.update(*args, **kwargs) create_entry = _MinEntry #defaults to a min-pq @classmethod def maxpq(cls, *args, **kwargs): pq = cls() pq.create_entry = _MaxEntry pq.__init__(*args, **kwargs) return pq def __len__(self): return len(self._heap) def __iter__(self): for entry in self._heap: yield entry.dkey def __getitem__(self, dkey): return self._heap[self._position[dkey]].pkey def __setitem__(self, dkey, pkey): heap = self._heap position = self._position try: pos = position[dkey] except KeyError: # Add a new entry: # put the new entry at the end and let it bubble up pos = len(self._heap) heap.append(self.create_entry(dkey, pkey)) position[dkey] = pos self._swim(pos) else: # Update an existing entry: # bubble up or down depending on pkeys of parent and children heap[pos].pkey = pkey parent_pos = (pos - 1) >> 1 child_pos = 2*pos + 1 if parent_pos > -1 and heap[pos] < heap[parent_pos]: self._swim(pos) elif child_pos < len(heap): other_pos = child_pos + 1 if other_pos < len(heap) and not heap[child_pos] < heap[other_pos]: child_pos = other_pos if heap[child_pos] < heap[pos]: self._sink(pos) def __delitem__(self, dkey): heap = self._heap position = self._position pos = position.pop(dkey) entry_to_delete = heap[pos] # Take the very last entry and place it in the vacated spot. Let it # sink or swim until it reaches its new resting place. end = heap.pop(-1) if end is not entry_to_delete: heap[pos] = end position[end.dkey] = pos parent_pos = (pos - 1) >> 1 child_pos = 2*pos + 1 if parent_pos > -1 and heap[pos] < heap[parent_pos]: self._swim(pos) elif child_pos < len(heap): other_pos = child_pos + 1 if other_pos < len(heap) and not heap[child_pos] < heap[other_pos]: child_pos = other_pos if heap[child_pos] < heap[pos]: self._sink(pos) del entry_to_delete def peek(self): try: entry = self._heap[0] except IndexError: raise KeyError return entry.dkey, entry.pkey def popitem(self): heap = self._heap position = self._position try: end = heap.pop(-1) except IndexError: raise KeyError if heap: entry = heap[0] heap[0] = end position[end.dkey] = 0 self._sink(0) else: entry = end del position[entry.dkey] return entry.dkey, entry.pkey def iteritems(self): # destructive heapsort iterator try: while True: yield self.popitem() except KeyError: return def _sink(self, top=0): # "Sink-to-the-bottom-then-swim" algorithm (Floyd, 1964) # Tends to reduce the number of comparisons when inserting "heavy" items # at the top, e.g. during a heap pop heap = self._heap position = self._position # Grab the top entry pos = top entry = heap[pos] # Sift up a chain of child nodes child_pos = 2*pos + 1 while child_pos < len(heap): # choose the smaller child other_pos = child_pos + 1 if other_pos < len(heap) and not heap[child_pos] < heap[other_pos]: child_pos = other_pos child_entry = heap[child_pos] # move it up one level heap[pos] = child_entry position[child_entry.dkey] = pos # next level pos = child_pos child_pos = 2*pos + 1 # We are left with a "vacant" leaf. Put our entry there and let it swim # until it reaches its new resting place. heap[pos] = entry position[entry.dkey] = pos self._swim(pos, top) def _swim(self, pos, top=0): heap = self._heap position = self._position # Grab the entry from its place entry = heap[pos] # Sift parents down until we find a place where the entry fits. while pos > top: parent_pos = (pos - 1) >> 1 parent_entry = heap[parent_pos] if not entry < parent_entry: break heap[pos] = parent_entry position[parent_entry.dkey] = pos pos = parent_pos # Put entry in its new place heap[pos] = entry position[entry.dkey] = pos