Created
September 8, 2017 09:47
-
-
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.
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
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