1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
|
# ==============================================================================
# Copyright (C) 2019 - Philip Paquette, Steven Bocco
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
# details.
#
# You should have received a copy of the GNU Affero General Public License along
# with this program. If not, see <https://www.gnu.org/licenses/>.
# ==============================================================================
""" Priority Dict implementation """
import heapq
# ------------------------------------------------
# Adapted from (2018/03/14s):
# https://docs.python.org/3.6/library/heapq.html#priority-queue-implementation-notes
# Unlicensed
class PriorityDict(dict):
""" Priority Dictionary Implementation """
def __init__(self, **kwargs):
""" Initialize the priority queue.
:param kwargs: (optional) initial values for priority queue.
"""
self.__heap = [] # Heap for entries. An entry is a triple (priority value, key, valid entry flag (boolean)).
# Dict itself maps key to entries. We override some dict methods (see __getitem__() below)
# to always return priority value instead of entry as dict value.
dict.__init__(self)
for key, value in kwargs.items():
self[key] = value
def __setitem__(self, key, val):
""" Sets a key with his associated priority
:param key: The key to set in the dictionary
:param val: The priority to associate with the key
:return: None
"""
if key in self:
del self[key]
# Create entry with val, key and a boolean indicating that entry is valid (True).
entry = [val, key, True]
dict.__setitem__(self, key, entry)
heapq.heappush(self.__heap, entry)
def __delitem__(self, key):
""" Removes key from dict and marks associated heap entry as invalid (False).
Raises KeyError if not found.
"""
entry = self.pop(key)
entry[-1] = False
def __getitem__(self, key):
""" Returns priority value associated to key. Raises KeyError if key not found. """
return dict.__getitem__(self, key)[0]
def __iter__(self):
""" Iterator over all keys based on their priority. """
def iterfn():
""" Iterator """
copy_of_self = self.copy()
while copy_of_self:
_, key = copy_of_self.smallest()
del copy_of_self[key]
yield key
return iterfn()
def smallest(self):
""" Finds the smallest item in the priority dict
:return: A tuple of (priority, key) for the item with the smallest priority, or None if dict is empty.
"""
while self.__heap and not self.__heap[0][-1]:
heapq.heappop(self.__heap)
return self.__heap[0][:2] if self.__heap else None
def setdefault(self, key, d=None):
""" Sets a default for a given key """
if key not in self:
self[key] = d
return self[key]
def copy(self):
""" Return a copy of this priority dict.
:rtype: PriorityDict
"""
return PriorityDict(**{key: entry[0] for key, entry in dict.items(self)})
def keys(self):
""" Make sure keys() iterates on keys based on their priority. """
return self.__iter__()
def values(self):
""" Makes sure values() iterates on priority values (instead of heap entries) from smallest to highest. """
return (self[k] for k in self)
def items(self):
""" Makes sure items() values are priority values instead of heap entries. """
return ((key, self[key]) for key in self)
|