sortedcontainers
介绍
本篇文章将主要介绍sortedcontainers中各个容器的实现方式。
第三方库地址：https://github.com/afthill/sorted_containers
python中含有大量的容器类型，比如list、set、dict等，但这些数据结构的有序版本却没有在标准库中实现。而在某些时候，可能需要一种能够动态维持有序的数据结构，在有序的基础上，使用二分查找来加快查询速度。
比如力扣的1847题（https://leetcode.cn/problems/closest-room/description/），官方题解就用到了python的第三方库sortedcontainers中的SortedList。
class Event:
"""
op: 事件的类型，0 表示房间，1 表示询问
size: 房间的 size 或者询问的 minSize
idx: 房间的 roomId 或者询问的 preferred
origin: 房间在数组 room 中的原始编号或者询问在数组 queries 中的原始编号
"""
def __init__(self, op: int, size: int, idx: int, origin: int):
self.op = op
self.size = size
self.idx = idx
self.origin = origin
"""
自定义比较函数，按照事件的 size 降序排序
如果 size 相同，优先考虑房间
"""
def __lt__(self, other: "Event") -> bool:
return self.size > other.size or (self.size == other.size and self.op < other.op)
class Solution:
def closestRoom(self, rooms: List[List[int]], queries: List[List[int]]) -> List[int]:
n = len(queries)
events = list()
for i, (roomId, size) in enumerate(rooms):
# 房间事件
events.append(Event(0, size, roomId, i))
for i, (minSize, preferred) in enumerate(queries):
# 询问事件
events.append(Event(1, preferred, minSize, i))
events.sort()
ans = [-1] * n
# 存储房间 roomId 的有序集合
# 需要导入 sortedcontainers 库
import sortedcontainers
valid = sortedcontainers.SortedList()
for event in events:
if event.op == 0:
# 房间事件，将 roomId 加入有序集合
valid.add(event.idx)
else:
# 询问事件
dist = float("inf")
# 查找最小的大于等于 preferred 的元素
x = valid.bisect_left(event.idx)
if x != len(valid) and valid[x] - event.idx < dist:
dist = valid[x] - event.idx
ans[event.origin] = valid[x]
if x != 0:
# 查找最大的严格小于 preferred 的元素
x -= 1
if event.idx - valid[x] <= dist:
dist = event.idx - valid[x]
ans[event.origin] = valid[x]
return ans
因此出于好奇，便去看了下这个库的实现。
SortedList
SortedList为sortedcontainers库中最重要的组件，其他的组件基本都是在其基础上设计的。
SortedList分为两种版本，一种比较元素时使用元素自身进行排序的，命名为SortedList，一种是按照传入的key函数进行排序的版本，命名为SortedKeyList。
class SortedList(MutableSequence)
class SortedKeyList(SortedList)
当初始化时，可以通过调用父类SortedList时是否传入key来自动生成SortedList或者SortedKeyList。
sl = SortedList([3, 4, 1, 2, 5])
print(sl)
print(type(sl))
skl = SortedList([3, 4, 1, 2, 5], key=neg)
print(skl)
print(type(skl))
"""
SortedList([1, 2, 3, 4, 5])
<class 'sortedcontainers.sortedlist.SortedList'>
SortedKeyList([5, 4, 3, 2, 1], key=<built-in function neg>)
<class 'sortedcontainers.sortedlist.SortedKeyList'>
"""
根据是否传入key生成不同类型的对象，主要是通过__new__方法实现的：
class SortedList(MutableSequence):
def __new__(cls, iterable=None, key=None):
if key is None:
return object.__new__(cls)
else:
if cls is SortedList:
return object.__new__(SortedKeyList)
else:
raise TypeError('inherit SortedKeyList for key argument')
__new__方法会在__init__方法之前进行，cls表示当前类，如果没有传入key，那么就直接构造当前类即可。如果传入了key，则会构造为SortedKeyList。
一个SortedList包含的主要成员为：
class SortedList(MutableSequence):
def __init__(self, iterable=None, key=None):
assert key is None
self._len = 0
self._load = self.DEFAULT_LOAD_FACTOR
self._lists = []
self._maxes = []
self._index = []
self._offset = 0
if iterable is not None:
self._update(iterable)
_len存储的是SortedList中元素的个数。
_list存储的是子列表，在SortedList将原先的有序长链表划分为许多子列表。
_maxes存储的是每个子列表的最大值。
_load表示装载因子，当子列表大于一定的值时发生分裂，或者小于一定值时发生合并。
_index是一个完全二叉树，记录第i到第j个子列表的元素数目（该节点为根的子树，最左叶子节点为第i个子列表，最右叶子节点为第j个子列表）。
_offset表示第一个叶子节点所在的下标，_index的完全二叉树是按照数组存储的。
一个SortedList的示意图如下：
如图所示的SortedList存储了1-14的元素，每个子列表都按照顺序村粗元素，且后面的列表中的元素均大于前面的列表。_maxes存储着每个子列表中的最大元素。_index的最底层为_lists的元素个数，并且以其为叶子节点，构建完全二叉树，每个节点的数字为子节点的数字之和。_offset则表示了在_index的数组中，叶子节点开始的位置（也是_lists[0]的长度存储的节点下标）。
SortedList分层存储信息，当查找元素时，根据_maxes找到元素可能在的桶，再在有序的桶中进行查找。默认的_load为1000时，作者测试此时SortedList对于几十到上千万的元素的操作性能都比较好。
单个元素的插入：
class SortedList(MutableSequence):
def add(self, value):
_lists = self._lists
_maxes = self._maxes
if _maxes:
pos = bisect_right(_maxes, value)
if pos == len(_maxes):
pos -= 1
_lists[pos].append(value)
_maxes[pos] = value
else:
insort(_lists[pos], value)
self._expand(pos)
else:
_lists.append([value])
_maxes.append(value)
self._len += 1
在插入前，会先在_maxes层进行查找，如果当前元素比列表内所有元素都大，则会插入到最后一个子列表中，否则就插入到刚好比_maxes小的那个子列表中。在插入之后再对子列表是否需要分裂进行判断。
多个元素的插入：
class SortedList(MutableSequence):
def update(self, iterable):
_lists = self._lists
_maxes = self._maxes
values = sorted(iterable)
if _maxes:
if len(values) * 4 >= self._len:
_lists.append(values)
values = reduce(iadd, _lists, [])
values.sort()
self._clear()
else:
_add = self.add
for val in values:
_add(val)
return
_load = self._load
_lists.extend(values[pos:(pos + _load)]
for pos in range(0, len(values), _load))
_maxes.extend(sublist[-1] for sublist in _lists)
self._len = len(values)
del self._index[:]
如果插入的元素较少，则会对每个元素都调用单个的add函数。如果插入的元素较多，则会将列表中的元素和待插入的元素进行排序，按照每个子列表为_load的大小对排序后的元素进行划分，并且更新相应的数据。
列表元素的删除：
class SortedList(MutableSequence):
def discard(self, value):
_maxes = self._maxes
if not _maxes:
return
pos = bisect_left(_maxes, value)
if pos == len(_maxes):
return
_lists = self._lists
idx = bisect_left(_lists[pos], value)
if _lists[pos][idx] == value:
self._delete(pos, idx)
找到对应的子列表后，再在子列表中删除对应坐标的元素。如果删除后，子列表过短，则会发生合并。
子列表的分裂：
当单个的子列表中存储了太多元素时，会发生子列表的分裂，使得分裂后子列表的元素均在阈值以下。
class SortedList(MutableSequence):
def _expand(self, pos):
_load = self._load
_lists = self._lists
_index = self._index
if len(_lists[pos]) > (_load << 1):
_maxes = self._maxes
_lists_pos = _lists[pos]
half = _lists_pos[_load:]
del _lists_pos[_load:]
_maxes[pos] = _lists_pos[-1]
_lists.insert(pos + 1, half)
_maxes.insert(pos + 1, half[-1])
del _index[:]
else:
if _index:
child = self._offset + pos
while child:
_index[child] += 1
child = (child - 1) >> 1
_index[0] += 1
如果子列表的长度超过阈值，则分出一段元素作为新列表，然后插入到_lists中，并且更新_maxes，同时删除索引_index。否则的话根据需要更新_index，更新_index时从叶子节点向上冒泡，依次更新。
合并的操作和分裂差不多，不多赘述。
指定元素的下标：
_index的一个重要作用就是辅助获取元素的下标。
class SortedList(MutableSequence):
def _loc(self, pos, idx):
if not pos:
return idx
_index = self._index
if not _index:
self._build_index()
total = 0
pos += self._offset
while pos:
if not pos & 1:
total += _index[pos - 1]
pos = (pos - 1) >> 1
return total + idx
上述函数会获取第pos个子列表中第idx个元素在全局中的下标。_loc会计算pos子列表前所有的子列表的长度之和，记为total，最终返回即total + idx。
_index由于是树状结构，大大减少了统计前面所有子列表长度和所需要的时间。算法从该子列表所对应的叶子节点开始，若为右节点，则加上左节点元素的个数。
例如对于_index = [14 5 9 3 2 4 5]，_offset=3，pos=2，idx=3举例，树可以参考上面的SortedList示意图。4为左孩子，因此total = 0；9为右孩子，因此total += 5，到达根节点结束。因此第2个子列表前的元素长度和为5（2+3）。所以第2个子列表中的第3个元素在全局的位置为8，也就是上图的元素8.
其他的函数感兴趣的朋友可以去阅读下源码。
SortedKeyList
SortedKeyList相较于SortedList，额外使用key作为元素的键值。其基本的数据结构为：
class SortedKeyList(SortedList):
def __init__(self, iterable=None, key=identity):
self._key = key
self._len = 0
self._load = self.DEFAULT_LOAD_FACTOR
self._lists = []
self._keys = []
self._maxes = []
self._index = []
self._offset = 0
if iterable is not None:
self._update(iterable)
除了使用_key来键值的获取方法外，还使用了_keys来存储每个元素使用_key获取的键值，可以将SortedList看作key和value一样的SortedKeyList。
以add函数为例，比较二者的不同：
class SortedKeyList(SortedList):
def add(self, value):
_lists = self._lists
_keys = self._keys
_maxes = self._maxes
key = self._key(value)
if _maxes:
pos = bisect_right(_maxes, key)
if pos == len(_maxes):
pos -= 1
_lists[pos].append(value)
_keys[pos].append(key)
_maxes[pos] = key
else:
idx = bisect_right(_keys[pos], key)
_lists[pos].insert(idx, value)
_keys[pos].insert(idx, key)
self._expand(pos)
else:
_lists.append([value])
_keys.append([key])
_maxes.append(key)
self._len += 1
例如在查找第pos个子列表和子列表中的下标idx时，都使用的key进行查找，且同时维护了_lists和_keys两个子列表数组。
SortedSet
SortedSet的实现机理并不复杂，其在内部除了python内置的set外，还有上述提到的SortedList。
class SortedSet(MutableSet, Sequence):
def __init__(self, iterable=None, key=None):
self._key = key
if not hasattr(self, '_set'):
self._set = set()
self._list = SortedList(self._set, key=key)
# 集合操作函数和有序列表操作函数
# 。。。
因此SortedSet既可以进行集合操作，也可进行列表操作。
ss1 = SortedSet([3, 4, 1, 2, 5], key=neg)
print(ss1)
ss2 = SortedSet([3, 6, 9], key=neg)
print(ss2)
print(ss1 & ss2)
print(ss1[:2])
"""
SortedSet([5, 4, 3, 2, 1], key=<built-in function neg>)
SortedSet([9, 6, 3], key=<built-in function neg>)
SortedSet([3], key=<built-in function neg>)
[5, 4]
"""
由于实现方式没有什么特殊的地方，主要还是靠set和SortedList互相补充，因此不作分析。
SortedDict
SortedDict和SortedSet类似，有序主要依靠内置的SortedList。
class SortedDict(dict):
def __init__(self, *args, **kwargs):
if args and (args[0] is None or callable(args[0])):
_key = self._key = args[0]
args = args[1:]
else:
_key = self._key = None
self._list = SortedList(key=_key)
# 其他有序列表的函数
# 。。。
SortedDict是直接继承自dict的，因此不需要像SortedSet一样，再加个set作为成员。
因此除了字典的操作外，SortedDict还可以进行按照key排序，第idx个item等有序操作。
sd = SortedDict({'a': 1, 'b': 2, 'c': 3})
print(sd)
print(sd.peekitem(1))
for k, v in sd.items():
print(k, v)
"""
SortedDict({'a': 1, 'b': 2, 'c': 3})
('b', 2)
a 1
b 2
c 3
"""
由于实现方式没有什么特殊的地方，因此不做分析。
总结
sortedcontainers的有序主要依靠SortedList实现，而SortedList则通过将较长的有序序列分割为较小的序列，减少了更新元素带来的性能损失。其实现方式也和其他有序容器比如跳表和红黑树等有一定的差异，在实战中应该根据项目需求，选择合适的容器。