Created
June 13, 2018 01:27
-
-
Save whiledoing/da88963b3f69cf44a2cd6b6654837dc0 to your computer and use it in GitHub Desktop.
[python-heap-with-updated-dict] python heap dict which can update heap in O(1), use redundant element for laze evaluation #python #datastructure
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# copyed from http://code.activestate.com/recipes/522995-priority-dict-a-priority-queue-with-updatable-prio/ | |
from heapq import heapify, heappush, heappop | |
class priority_dict(dict): | |
"""Dictionary that can be used as a priority queue. | |
Keys of the dictionary are items to be put into the queue, and values | |
are their respective priorities. All dictionary methods work as expected. | |
The advantage over a standard heapq-based priority queue is | |
that priorities of items can be efficiently updated (amortized O(1)) | |
using code as 'thedict[item] = new_priority.' | |
The 'smallest' method can be used to return the object with lowest | |
priority, and 'pop_smallest' also removes it. | |
The 'sorted_iter' method provides a destructive sorted iterator. | |
""" | |
def __init__(self, *args, **kwargs): | |
super(priority_dict, self).__init__(*args, **kwargs) | |
self._rebuild_heap() | |
def _rebuild_heap(self): | |
self._heap = [(v, k) for k, v in self.iteritems()] | |
heapify(self._heap) | |
def smallest(self): | |
"""Return the item with the lowest priority. | |
Raises IndexError if the object is empty. | |
""" | |
heap = self._heap | |
v, k = heap[0] | |
while k not in self or self[k] != v: | |
heappop(heap) | |
v, k = heap[0] | |
return k | |
def pop_smallest(self): | |
"""Return the item with the lowest priority and remove it. | |
Raises IndexError if the object is empty. | |
""" | |
heap = self._heap | |
v, k = heappop(heap) | |
while k not in self or self[k] != v: | |
v, k = heappop(heap) | |
del self[k] | |
return k | |
def __setitem__(self, key, val): | |
# We are not going to remove the previous value from the heap, | |
# since this would have a cost O(n). | |
super(priority_dict, self).__setitem__(key, val) | |
if len(self._heap) < 2 * len(self): | |
heappush(self._heap, (val, key)) | |
else: | |
# When the heap grows larger than 2 * len(self), we rebuild it | |
# from scratch to avoid wasting too much memory. | |
self._rebuild_heap() | |
def setdefault(self, key, val): | |
if key not in self: | |
self[key] = val | |
return val | |
return self[key] | |
def update(self, *args, **kwargs): | |
# Reimplementing dict.update is tricky -- see e.g. | |
# http://mail.python.org/pipermail/python-ideas/2007-May/000744.html | |
# We just rebuild the heap from scratch after passing to super. | |
super(priority_dict, self).update(*args, **kwargs) | |
self._rebuild_heap() | |
def sorted_iter(self): | |
"""Sorted iterator of the priority dictionary items. | |
Beware: this will destroy elements as they are returned. | |
""" | |
while self: | |
yield self.pop_smallest() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment