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