File size: 2,245 Bytes
61517de |
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 |
import heapq
class SortedDict(dict):
def __init__(self, sort_func=lambda k, v: k, init_dict=None, reverse=False):
if init_dict is None:
init_dict = []
if isinstance(init_dict, dict):
init_dict = init_dict.items()
self.sort_func = sort_func
self.sorted_keys = None
self.reverse = reverse
self.heap = []
for k, v in init_dict:
self[k] = v
def __setitem__(self, key, value):
if key in self:
super().__setitem__(key, value)
for i, (priority, k) in enumerate(self.heap):
if k == key:
self.heap[i] = (self.sort_func(key, value), key)
heapq.heapify(self.heap)
break
self.sorted_keys = None
else:
super().__setitem__(key, value)
heapq.heappush(self.heap, (self.sort_func(key, value), key))
self.sorted_keys = None
def __delitem__(self, key):
super().__delitem__(key)
for i, (priority, k) in enumerate(self.heap):
if k == key:
del self.heap[i]
heapq.heapify(self.heap)
break
self.sorted_keys = None
def keys(self):
if self.sorted_keys is None:
self.sorted_keys = [k for _, k in sorted(self.heap, reverse=self.reverse)]
return self.sorted_keys
def items(self):
if self.sorted_keys is None:
self.sorted_keys = [k for _, k in sorted(self.heap, reverse=self.reverse)]
sorted_items = [(k, self[k]) for k in self.sorted_keys]
return sorted_items
def _update_heap(self, key):
for i, (priority, k) in enumerate(self.heap):
if k == key:
new_priority = self.sort_func(key, self[key])
if new_priority != priority:
self.heap[i] = (new_priority, key)
heapq.heapify(self.heap)
self.sorted_keys = None
break
def __iter__(self):
return iter(self.keys())
def __repr__(self):
return f"{type(self).__name__}({dict(self)}, sort_func={self.sort_func.__name__}, reverse={self.reverse})"
|