|
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})" |
|
|