Skip to content

Instantly share code, notes, and snippets.

@qfgaohao
Created September 8, 2017 09:47
Show Gist options
  • Save qfgaohao/bddfcdec5a0f1b6cc845b01708493145 to your computer and use it in GitHub Desktop.
Save qfgaohao/bddfcdec5a0f1b6cc845b01708493145 to your computer and use it in GitHub Desktop.
A structure combines priority queue and dict. You can pop out the biggest item from it like in queue. You can also set and get item like on a dict, while the order in the priority queue is maintained.
class Pair:
def __init__(self, key, value):
self.key = key
self.value = value
def __str__(self):
return "({}, {})".format(str(self.key), str(self.value))
def __repr__(self):
return self.__str__()
class PriorityMap:
"""It combines a priority queue and a dict.
The queue is a heap ordered by value. The biggest value is in the begining of the queue.
"""
def __init__(self, items=[]):
self.map = {} # {key: index in heap}
self.heap = [] # [value]
for i, (key, value) in enumerate(items):
self.map[key] = i
self.heap.append(Pair(key, value))
self._heapify()
def replace(self, key, value):
if key not in self.map:
raise KeyError('The key is not in the map.')
pair = Pair(key, value)
i = self.map[key]
rt = self.heap[i]
self.heap[i] = pair
if pair.value < rt.value:
self._sift_down(i)
elif pair.value > rt.value:
self._sift_up(i)
return rt.key, rt.value
def pop(self):
if not self.heap:
raise ValueError('The priority map is empty.')
rt = self.heap[0]
del self.map[rt.key]
self.heap[0] = self.heap[-1]
self.map[self.heap[0].key] = 0
self.heap.pop() # remove last
self._sift_down(0)
return rt.key, rt.value
def __setitem__(self, key, value):
if key in self.map:
self.replace(key, value)
else:
self.append(key, value)
def __getitem__(self, key):
if key not in self.map:
raise KeyError('The key is not in the map.')
return self.heap[self.map[key]].value
def __delitem__(self, key):
if key not in self.map:
raise KeyError('The key is not in the map')
i = self.map[key]
self.heap[i] = self.heap[-1]
self.heap.pop() # list::pop
del self.map[key]
pair = self.heap[i]
parent = (i - 1) // 2
p = self.heap[parent]
if pair.value < p.value:
self._sift_down(i)
elif pair.value > p.value:
self._sift_up(i)
def append(self, key, value):
if key in self.map:
raise KeyError('The key is already in the priority map.')
self.map[key] = len(self.map)
self.heap.append(Pair(key, value))
self._sift_up(len(self.heap) - 1)
def __len__(self):
return len(self.heap)
def _heapify(self):
start = (len(self.heap) - 1) // 2
while start >= 0:
self._sift_down(start)
start -= 1
def _sift_down(self, start):
if 2 * start + 1 >= len(self.heap):
return
left = 2 * start + 1
big = left
if (2 * start + 2 < len(self.heap)):
right = 2 * start + 2
if self.heap[right].value > self.heap[big].value:
big = right
if self.heap[start].value < self.heap[big].value:
self._swap(start, big)
self._sift_down(big)
def _sift_up(self, start):
if start <= 0:
return
parent = (start - 1) // 2
if self.heap[start].value > self.heap[parent].value:
self._swap(start, parent)
self._sift_up(parent)
def _swap(self, i, j):
self.map[self.heap[i].key] = j
self.map[self.heap[j].key] = i
p = self.heap[i]
self.heap[i] = self.heap[j]
self.heap[j] = p
pair_list2tuple_list = lambda pairs: [(p.key, p.value) for p in pairs]
m = PriorityMap([('a', 1), ('b', 3), ('c', -2), ('d', 5)])
assert pair_list2tuple_list(m.heap) == [('d', 5), ('b', 3), ('c', -2), ('a', 1)]
assert m.map == {'a': 3, 'b': 1, 'c': 2, 'd': 0}
m.append('e', 4)
assert pair_list2tuple_list(m.heap) == [('d', 5), ('e', 4), ('c', -2), ('a', 1), ('b', 3)]
assert m.map == {'a': 3, 'b': 4, 'c': 2, 'd': 0, 'e': 1}
m['f'] = 8
assert pair_list2tuple_list(m.heap) == [('f', 8), ('e', 4), ('d', 5), ('a', 1), ('b', 3), ('c', -2)]
assert m.map == {'a': 3, 'b': 4, 'c': 5, 'd': 2, 'e': 1, 'f': 0}
m.replace('e', 9)
assert pair_list2tuple_list(m.heap) == [('e', 9), ('f', 8), ('d', 5), ('a', 1), ('b', 3), ('c', -2)]
assert m.map == {'a': 3, 'b': 4, 'c': 5, 'd': 2, 'e': 0, 'f': 1}
assert m.pop() == ('e', 9)
assert pair_list2tuple_list(m.heap) == [('f', 8), ('b', 3), ('d', 5), ('a', 1), ('c', -2)]
assert m.map == {'a': 3, 'b': 1, 'c': 4, 'd': 2, 'f': 0}
del m['d']
assert pair_list2tuple_list(m.heap) == [('f', 8), ('b', 3), ('c', -2), ('a', 1)]
assert m.map == {'a': 3, 'b': 1, 'c': 4, 'f': 0}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment