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
Diff to Previous Revision
--- revision 3 2013-08-18 16:17:54
+++ revision 4 2013-08-25 00:23:12
@@ -1,26 +1,20 @@
-"""
-Entry classes customize the ranking behavior of the priority queue. The heapify
-algorithms of PQDict use the "<" comparator to compare entries, so entry classes
-must implement __lt__.
+from collections import MutableMapping
-"""
-from collections import Mapping, MutableMapping
-
-class MinPQDEntry(object):
+class _MinEntry(object):
"""
- Entries for a PQDict backed by a min-heap.
+ Mutable entries for a Min-PQ dictionary.
"""
def __init__(self, dkey, pkey):
- self.dkey = dkey
- self.pkey = pkey
+ self.dkey = dkey #dictionary key
+ self.pkey = pkey #priority key
def __lt__(self, other):
return self.pkey < other.pkey
-class MaxPQDEntry(object):
+class _MaxEntry(object):
"""
- Entries for a PQDict backed by a max-heap.
+ Mutable entries for a Max-PQ dictionary.
"""
def __init__(self, dkey, pkey):
@@ -31,81 +25,44 @@
return self.pkey > other.pkey
class PQDict(MutableMapping):
- create_entry = MinPQDEntry
+ 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 = MaxPQDEntry
+ pq.create_entry = _MaxEntry
pq.__init__(*args, **kwargs)
return pq
- def __init__(self, *args, **kwargs):
- """
- Same input signature as dict:
- Accepts a sequence/iterator of (dkey, pkey) pairs.
- Accepts named arguments or an unpacked dictionary.
- Also accepts a single mapping object to convert it to a pqdict.
-
- """
- if len(args) > 1:
- raise TypeError
- self._heap = []
- self._posfinder = {}
- pos = 0
-
- if args:
- if isinstance(args[0], Mapping):
- seq = args[0].items()
- else:
- seq = args[0]
- try:
- for dkey, pkey in seq:
- entry = self.create_entry(dkey, pkey)
- self._heap.append(entry)
- self._posfinder[dkey] = pos
- pos += 1
- except TypeError:
- raise ValueError
- if kwargs:
- for dkey, pkey in kwargs.items():
- entry = self.create_entry(dkey, pkey)
- self._heap.append(entry)
- self._posfinder[dkey] = pos
- pos += 1
- self._heapify()
+ def __len__(self):
+ return len(self._heap)
def __iter__(self):
- """
- Destructive heapsort iterator over items, ordered by priority key.
-
- """
- try:
- while True:
- yield self.popitem()
- except KeyError:
- return
-
- def __len__(self):
- return len(self._posfinder)
+ for entry in self._heap:
+ yield entry.dkey
def __getitem__(self, dkey):
- return self._heap[self._posfinder[dkey]].pkey
+ return self._heap[self._position[dkey]].pkey
def __setitem__(self, dkey, pkey):
heap = self._heap
- finder = self._posfinder
+ position = self._position
+
try:
- pos = finder[dkey]
+ pos = position[dkey]
except KeyError:
- # add new entry:
+ # Add a new entry:
# put the new entry at the end and let it bubble up
- n = len(self._heap)
- self._heap.append(self.create_entry(dkey, pkey))
- self._posfinder[dkey] = n
- self._swim(n)
+ pos = len(self._heap)
+ heap.append(self.create_entry(dkey, pkey))
+ position[dkey] = pos
+ self._swim(pos)
else:
- # update existing entry:
+ # Update an existing entry:
# bubble up or down depending on pkeys of parent and children
heap[pos].pkey = pkey
parent_pos = (pos - 1) >> 1
@@ -113,57 +70,36 @@
if parent_pos > -1 and heap[pos] < heap[parent_pos]:
self._swim(pos)
elif child_pos < len(heap):
- child2_pos = child_pos + 1
- if (child2_pos < len(heap)
- and not heap[child_pos] < heap[child2_pos]):
- child_pos = child2_pos
+ 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
- finder = self._posfinder
- try:
- pos = finder.pop(dkey)
- except KeyError:
- raise
- else:
- # Take the very last entry and place it in the vacated spot. Let it
- # sink or swim until it reaches its new resting place.
- entry_to_delete = heap[pos]
- last_entry = heap.pop(-1)
- if entry_to_delete is not last_entry:
- heap[pos] = last_entry
- finder[last_entry.dkey] = pos
+ position = self._position
- 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):
- child2_pos = child_pos + 1
- if (child2_pos < len(heap) and
- not heap[child_pos] < heap[child2_pos]):
- child_pos = child2_pos
- if heap[child_pos] < heap[pos]:
- self._sink(pos)
- del entry_to_delete
+ pos = position.pop(dkey)
+ entry_to_delete = heap[pos]
- def popitem(self):
- try:
- last = self._heap.pop(-1)
- except IndexError:
- raise KeyError
- else:
- if self._heap:
- entry = self._heap[0]
- self._heap[0] = last
- self._posfinder[last.dkey] = 0
- self._sink(0)
- else:
- entry = last
- del self._posfinder[entry.dkey]
- return entry.dkey, entry.pkey
+ # 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:
@@ -172,61 +108,78 @@
raise KeyError
return entry.dkey, entry.pkey
- def _heapify(self):
- n = len(self._heap)
- for pos in reversed(range(n//2)):
- self._sink(pos)
+ 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
- finder = self._posfinder
+ 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.
- right_pos = child_pos + 1
- if right_pos < len(heap) and not heap[child_pos] < heap[right_pos]:
- child_pos = right_pos
+ # 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.
+ # move it up one level
heap[pos] = child_entry
- finder[child_entry.dkey] = pos
+ 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
- finder[entry.dkey] = pos
+ position[entry.dkey] = pos
self._swim(pos, top)
def _swim(self, pos, top=0):
heap = self._heap
- finder = self._posfinder
+ 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 entry < parent_entry:
- heap[pos] = parent_entry
- finder[parent_entry.dkey] = pos
- pos = parent_pos
- continue
- break
-
+ 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
- finder[entry.dkey] = pos
+ position[entry.dkey] = pos
