mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-01-19 01:21:01 +08:00
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
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})"
|