分享|SortedList 源码解读与优化

111
2026.10.30
2026.10.30
发布于 福建
1 设计思路
1.1 问题的起点
1.3 第二步:分块维护块内元素
1.4 第三步:树形结构维护块长
1.5 复杂度小结
2.1 类属性
2.2 线段树构建索引
2.3 下标换算
2.4 块的分裂
2.5 块的合并
2.6 索引惰性重建
2.7 按值操作
2.7.1 定位值
2.7.2 插入值
2.7.3 删除值
2.8 按下标操作
2.8.1 按下标访问
2.8.2 按下标删除
2.9 批量插入与删除
2.10 区间遍历
2.10.1 全量迭代
2.10.2 按下标区间迭代
2.10.3 按值区间迭代
2.11 其他细节
2.11.1 使用内置库
2.11.2 局部变量缓存
2.11.3 特判快路径
2.12 复杂度小结
3.1 树状数组构建索引
3.2 其余差异
4.1 构建索引
4.2 块长取值
4.3 小块合并
4.4 功能完善
4.4.1 返回值
4.4.2 区间迭代
4.4.3 区间删除
4.5 其他细节
4.6 复杂度小结
5 基准测试
5.1 测试标准
5.2 不同块长测试
5.3 不同版本测试
5.4 在线评测
6 总结
相关阅读
1 设计思路
1.1 问题的起点
Python 的标准库里没有平衡树。在算法竞赛里,C++ 有 std::set、std::multiset 以及 pbds 的 tree、ordered_set。为了对标它们,Python 需要一种至少能实现以下四个功能的数据结构:
  • 插入、删除一个值(允许重复)
  • 给定值、查询名次
  • 给定名次、查询值
  • 有序遍历区间
为什么不直接写平衡树?Treap、Splay、红黑树的实现都是几百行,代码量尚可接受。但真正的问题是常数因子。以红黑树为例,如果使用类表示红黑树节点,由于节点是逐个分配的、内存散落在堆上,物理连续性相当差;即使改用数组下标表示节点,各个节点的访问顺序也是随机的。沿着红黑树的树高访问和修改 O(log⁡n)O(\log n) 个节点,整体代价偏高。
因此,在 Python 里照搬 C++ 的数据结构形态是错的方向,必须换一种对解释器友好的内存布局。
1.2 第一步:list 维护有序
list 是连续内存,按下标访问为 O(1)O(1),bisect.insort 二分找到位置为 O(log⁡n)O(\log n),但有序表的插入和删除需要整体移动数组元素,平均代价 O(n)O(n),单次增删的用时随 nn 线性增长。
元素移动的底层是 C 的 memmove,在 CPython 上实测约 0.1ns/元素,比在 Python 里手写循环逐个移动快两个数量级,PyPy 下也在同一量级。也就是说,只要可以控制 list 的长度,就可以控制单次插入和删除操作的用时。
1.3 第二步:分块维护块内元素
于是可以用分块控制长度。把整个序列切成长为 LL 的块,块数约为 B=n/LB=n/L,每块内部有序,块与块之间整体有序。另存一个 maxes 数组记录每块的最大值。插入分两步:bisect_right(maxes, v) 找到候选块,复杂度 O(log⁡B)=O(log⁡(n/L))O(\log B)=O(\log (n/L));insort(block, v) 在块内插入,复杂度 O(log⁡L+L)O(\log L + L)。删除的定位方式相同,块内删除的移动量与插入一致。
设块内元素移动的单位成本为 c1c_1、索引比较的单位成本为 c2c_2,则单次插入的主要成本为 Lc1+log⁡(n/L)⋅c2L c_1 + \log (n/L)\cdot c_2。对 LL 求导取极值,则 L∗∼c2/c1L^* \sim c_2/c_1。代入实测的 c1≈0.1nsc_1 \approx 0.1\text{ns}、c2≈25nsc_2 \approx 25\text{ns},得 L∗≈250L^* \approx 250。因此块长取在几百的量级时,单次增删的代价是较低的。
块会随着插入变长,必须在超过阈值时分裂,否则会退化成一个超长的 list。解决方式也很简单:块长 >2L>2L 时在中间切开、成为两个新块,maxes 也要相应更新。分裂要移动 O(L)O(L) 个元素,但要连续插入 O(L)O(L) 次才会触发一次,摊还到每次插入仍是 O(1)O(1)。
1.4 第三步:树形结构维护块长
分块解决了按值插入的问题,但还需要解决按照名次访问元素的问题。按名次访问也分两步:给定全局下标 kk,找到它落在哪一块;给定块号 pp,找到块内对应的下标。这两步都依赖在块长数组上查找 kk。如果朴素维护块长数组,块号查询是 O(B)O(B) 的;反之如果维护块长数组的前缀和,块长度修改是 O(B)O(B) 的。这一代价比块内移动还高。
因此需要一种能高效维护块长的结构,树状数组和线段树都是可选方案。在块数不变的情况下,块长修改可以表示为数据结构的单点修改,求块号则可以表示为在前缀和上的二分,二者都是 O(log⁡B)=O(log⁡(n/L))O(\log B) = O(\log(n/L)) 的。
再考虑块数发生变化的情况。索引的叶子与块号一一绑定。块 pp 分裂成 pp、p+1p+1 之后,原来 p+1p+1 及其后的所有块号都要 +1+1。块号整体平移后,树的叶子和内部节点不再对应,无法靠单点修改修补。所以块数变化后必须按新的块长数组重新建一遍树。
由于每 O(L)O(L) 次插入才产生一次分裂,每次重建 O(B)=O(n/L)O(B) = O(n/L),摊还到每次插入是 O(n/L2)O(n/L^2)。n=106n = 10^6、L=250L = 250 时约为每次插入多付 16 次基本运算,可以接受。
1.5 复杂度小结
到这里,结构层面的设计结束,各操作的复杂度如下。LL 是唯一的自由参数:LL 小则块内移动少、块数多;LL 大则块内移动多、块数少。表中插入与删除含索引重建的摊还项。
操作
复杂度
插入 / 删除(摊还)
O(L+log⁡n+n/L2)O(L + \log n + n/L^2)
按值查询名次
O(log⁡n)O(\log n)
按名次查询值
O(log⁡(n/L))O(\log(n/L))
2 sortedcontainers 版 SortedList 源码解读
sortedcontainers(后文简称 std)是一个纯 Python 的有序容器库,提供 SortedList、SortedDict、SortedSet 三个类型,以及带 key 函数的 SortedKeyList。本文仅就 SortedList 作分析。
std 的文档里给出了不少设计理念和具体实现,与我们上文的分析相呼应。docs/implementation.rst 开篇给出两条观察:Python 的 list 很快,bisect.insort 也很快。单用一个 list 加 insort,长度超过一万才会明显退化;改用 list of lists 之后,插入和删除大多落在短块上,只有少数情况才需要增删整块。与树形结构相比,这种布局有四点内存优势:
多数插入和删除不涉及内存分配与释放,减轻垃圾回收器和内存系统的压力;元素指针密集排列,有利于缓存的利用;每个元素的内存额外开销实际上只是一个指针,而二叉树节点至少要再多存两个子指针;顺序遍历是处理器的强项,连续内存的迭代快于指针跳转。
在现代 CPU 上经过测试,顺序访问比随机访问快约 5 倍;数据依赖访问(每次地址依赖上一次结果)比随机访问慢约 5 倍。移动 1000 个连续元素,比走 30 层指针追踪更快。这是分块表在 Python 中胜过平衡树的根本原因,也是本文所有优化的前提。std 因此选择了类似两层 B+ 树的整体结构。
docs/performance.rst 还比较了 blist、bintrees、sortedmap、banyan 等 C 扩展实现,结论是纯 Python 的 std 在多数场景下不慢于这些 C 实现,落后时也基本在同一量级。
2.1 类属性
Python
class SortedList(MutableSequence):
    DEFAULT_LOAD_FACTOR = 1000

    def __init__(self, iterable=None, key=None):
        self._len = 0                          # 元素数量
        self._load = self.DEFAULT_LOAD_FACTOR  # 默认块长
        self._lists = []                       # 块内元素
        self._maxes = []                       # 块内最大值
        self._index = []                       # 块长索引
        self._offset = 0                       # 索引中第一个叶子节点的下标
        ...
std 选取了 DEFAULT_LOAD_FACTOR=1000 作为默认块长。文档给出的经验是,默认值适合长度不超过一千万的列表;更大时应把块长取为平均长度的平方根到立方根之间。后文会说明这个建议的来源。
std 记录了块最大值 _maxes 而非最小值 _mins,这样 bisect_left(_maxes, v) 落到的块必然包含 vv 的插入位置,而存最小值则要额外检查前一块的块尾、多做一次修正。
2.2 线段树构建索引
为了维护块长数组,std 选择了维护一个满二叉树扁平数组 _index,以及 _index 中第一个叶子节点的位置 _offset,_index 的后半部分是叶子节点、表示块长,前半部分则是子树中叶子节点的和。例如:
Python
_index = [14, 5, 9, 3, 2, 4, 5]   # 内部节点在前,叶子节点在后
_offset = 3                       # 第一个叶子节点的下标
对应:
PlainText
     14        ——内部节点
   /   \
  5     9      ——内部节点
 / \   / \
3   2 4   5    ——块长
索引构建见 _build_index。构建二叉树的关键细节是把第二层补零到 2 的幂,补零后除叶子节点外,每一层都成为满层,_offset 固定为 2k+1−12^{k+1}-1(内部节点数),下降循环只需判断 child < len(_index),不必再维护树的左右边界。这与自底向上的 ZKW 线段树相似,下文还会继续体现。
Python
def _build_index(self):
    row0 = list(map(len, self._lists))          # 叶子层:每块的块长
    ...                                         # 单块时提前返回
    head = iter(row0)
    tail = iter(head)                           # 同一个迭代器,zip 才能取到相邻对
    row1 = list(starmap(add, zip(head, tail)))  # 相邻块长两两求和,得到上一层
    if len(row0) & 1:                           # 奇数个叶子,末位直接上提
        row1.append(row0[-1])
    ...
    size = 2 ** (int(log(len(row1) - 1, 2)) + 1)  # 补零后的叶子层长度
    row1.extend(repeat(0, size - len(row1)))    # 补零到 2 的幂
    tree = [row0, row1]                         # 按层收集,从叶子到根
    while len(tree[-1]) > 1:
        head = iter(tree[-1])
        tail = iter(head)
        tree.append(list(starmap(add, zip(head, tail))))    # 线性建树
    reduce(iadd, reversed(tree), self._index)   # 从根到叶子拼成扁平数组
    self._offset = size * 2 - 1                 # 内部节点数即叶子起点
2.3 下标换算
块坐标转换为全局下标见 _loc。从叶子向根走,遇到右孩子就把左兄弟的子树和加进来,固定走满 log⁡B\log B 层。在计算之前检查 _index 是否为空,为空则调用 _build_index。
Python
def _loc(self, pos, idx):
    ...
    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。从根向下走,每层比较一次左孩子的和值。同样是 log⁡B\log B 层。
Python
def _pos(self, idx):
    ...
    if not _index:                       # 惰性重建
        self._build_index()

    pos = 0
    child = 1                           # 从根节点开始
    len_index = len(_index)
    while child < len_index:
        index_child = _index[child]
        if idx < index_child:
            pos = child                  # 走左子树
        else:
            idx -= index_child           # 扣掉左子树
            pos = child + 1              # 走右子树
        child = (pos << 1) + 1
    return pos - self._offset, idx
2.4 块的分裂
因插入元素导致块分裂见 _expand。块长超过 2L2L 时,在 LL 处切成两段,清空索引;否则块数不变,在线段树上增量更新。
Python
def _expand(self, pos):
    ...
    if len(_lists[pos]) > (_load << 1):    # 块长超过 2L,开始分裂
        half = _lists[pos][_load:]
        del _lists[pos][_load:]
        _lists.insert(pos + 1, half)
        _maxes[pos] = _lists[pos][-1]
        _maxes.insert(pos + 1, half[-1])
        del _index[:]                      # 索引作废,等下次读时重建
    else:
        if _index:                         # 索引已过期就跳过增量更新
            child = self._offset + pos
            while child:                   # 沿叶子到根 +1
                _index[child] += 1
                child = (child - 1) >> 1
            _index[0] += 1
2.5 块的合并
因删除元素导致块合并见 _delete。块长降到 L/2L/2 及以下时,与相邻块合并并清空索引。合并后的块可能超过 2L2L,再触发一次分裂。合并策略使除末块外的块长保持在 L/2L/2 到 2L2L 之间,块数稳定在 n/Ln/L 附近。
Python
def _delete(self, pos, idx):
    ...
    _lists_pos = _lists[pos]
    del _lists_pos[idx]
    self._len -= 1
    len_lists_pos = len(_lists_pos)

    if len_lists_pos > (self._load >> 1):  # 块长仍 > L/2
        _maxes[pos] = _lists_pos[-1]
        ...                                # 索引增量 -1,同 _expand
    elif len(_lists) > 1:                  # 块长 <= L/2,与邻块合并
        if not pos:
            pos += 1
        prev = pos - 1
        _lists[prev].extend(_lists[pos])   # 并入左邻块
        _maxes[prev] = _lists[prev][-1]
        del _lists[pos]
        del _maxes[pos]
        del _index[:]
        self._expand(prev)                 # 合并后可能超 2L,再分裂
    elif len_lists_pos:                    # 只剩一个块且未删空
        _maxes[pos] = _lists_pos[-1]
    else:                                  # 块被删空
        del _lists[pos]
        del _maxes[pos]
        del _index[:]
2.6 索引惰性重建
std 中索引要么与 _lists 的块长完全一致,要么整个作废,不存在部分更新的中间状态。代价是每次读索引都要先判断它是否可用;收益是块数变化时不需要立刻修改。
块数变化时,也就是 _expand 分裂出新块、_delete 合并或删空整块时,索引直接清空。重建的触发点是块数变化之后的第一次索引换算,而不是块数变化本身:_loc、_pos 在入口处检查 _index 是否为空,为空才调用 _build_index。这期间只有需要读索引的操作(bisect_left、__getitem__、pop 等)才会触发一次重建,只做 add、discard 不会。这样的摊还成本仍是 O(n/L2)O(n/L^2)。
2.7 按值操作
2.7.1 定位值
bisect_left 先在 _maxes 上二分定位块,再在块内二分定位下标,最后用 _loc 把 (块号, 块内下标) 换算成全局下标。bisect_right 只把两处 bisect_left 换成 bisect_right,其余相同。两级二分的比较次数是 O(log⁡(n/L)+log⁡L)=O(log⁡n)O(\log(n/L) + \log L) = O(\log n)。
Python
def bisect_left(self, value):
    ...
    pos = bisect_left(_maxes, value)            # 定位候选块
    idx = bisect_left(self._lists[pos], value)  # 块内定位
    return self._loc(pos, idx)                  # 换算成全局下标
__contains__ 只判等,不换算全局下标:
Python
def __contains__(self, value):
    ...
    pos = bisect_left(_maxes, value)            # 定位候选块
    idx = bisect_left(_lists[pos], value)       # 块内定位
    return _lists[pos][idx] == value
index 在判等之后用 _loc 求出首次出现的全局下标,另外还要处理 start、stop 区间。
Python
def index(self, value, start=None, stop=None):
    ...
    stop -= 1
    left = self._loc(pos_left, idx_left)       # 该值首次出现的全局下标
    if start <= left:
        if left <= stop:  # left 落在区间内
            return left
    elif start <= self._bisect_right(value) - 1:  # left 落在区间左侧,但 value 在 start 之后仍有出现
        return start
    raise ValueError(f'{value!r} is not in list')
count 分别求 value 的最左、最右插入位置,再相减。
2.7.2 插入值
add 分三步:定位块时用 bisect_right(_maxes, value),使相等的值排在已有元素之后;在块内插入;最后调用 _expand 维护块长,数量增加时增量更新索引,块长超过 2L2L 时分裂。
Python
def add(self, value):
    ...
    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
2.7.3 删除值
discard 的定位与判等和 __contains__ 相同,命中后再调用 _delete。remove 与之相同,只是未命中时抛 ValueError。
Python
def discard(self, value):
    ...
    pos = bisect_left(_maxes, value)
    if pos == len(_maxes):
        return
    idx = bisect_left(_lists[pos], value)
    if _lists[pos][idx] == value:
        self._delete(pos, idx)
2.8 按下标操作
2.8.1 按下标访问
整数下标先经过四次特判:首元素、末元素、首块、末块;都未命中再用 _pos 定位。
Python
def __getitem__(self, index):
    ...
    if self._len:
        if index == 0:                     # 首元素
            return _lists[0][0]
        elif index == -1:                  # 末元素
            return _lists[-1][-1]
    else:
        raise IndexError('list index out of range')
    if 0 <= index < len(_lists[0]):        # 首块
        return _lists[0][index]
    len_last = len(_lists[-1])
    if -len_last < index < 0:              # 末块
        return _lists[-1][len_last + index]
    pos, idx = self._pos(index)            # 快路径未命中,走索引定位
    return _lists[pos][idx]
std 的 __getitem__ 还支持切片,2.9 节的 __delitem__ 会用它进行区间删除。
2.8.2 按下标删除
pop 的快路径与 __getitem__ 一一对应,命中后直接 _delete;快路径未命中时用 _pos 定位再 _delete。
Python
def pop(self, index=-1):
    ...
    if index == 0:                        # 首元素
        val = _lists[0][0]
        self._delete(0, 0)
        return val
    if index == -1:                       # 末元素
        ...
    if 0 <= index < len(_lists[0]):       # 首块
        ...
    if -len(_lists[-1]) < index < 0:      # 末块
        ...
    pos, idx = self._pos(index)
    val = _lists[pos][idx]
    self._delete(pos, idx)
    return val
2.9 批量插入与删除
批量插入见 update。它不是单纯循环调用 add,而是比较新数据量与当前长度:新数据达到当前长度的 1/4 时,把所有元素取出排序、重新切块;否则逐个 add。新数据占比较大时,逐个 add 要付出 k×O(L+log⁡n)k \times O(L + \log n) 的代价,而摊平重排只需一次 O(n+k)O(n + k) 的切块、外加一次排序。sort 也不是从零排序:values 仅由两段有序段构成,Python 排序所用的 Timsort 会把这些现成段直接归并,实际代价接近 O(n+klog⁡k)O(n + k\log k)。这个阈值是权衡两种方式的经验值。
Python
def update(self, iterable):
    ...
    if len(values) * 4 >= self._len:      # 新数据 >= 现有量的 1/4
        _lists.append(values)
        values = reduce(iadd, _lists, []) # 摊平所有块
        values.sort()                     # 合并已有有序段
        self._clear()
    else:
        ...                               # 未达阈值,逐个 add
    ...                                   # 按 load 重新切块
批量删除见 __delitem__。它也用同样的思路,阈值是长度的 1/8,属于同一类经验阈值:删除量大时保留补集重建,是 O(n)O(n);逐个删是 m×O(L+log⁡n)m \times O(L + \log n)。
Python
def __delitem__(self, index):
    ...
    if self._len <= 8 * (stop - start):       # 删除量 >= 总量的 1/8
        values = self.__getitem__(slice(None, start))      # 取区间左侧
        if stop < self._len:
            values += self.__getitem__(slice(stop, None))  # 再拼区间右侧
        self._clear()
        return self._update(values)           # 保留补集,整体重建
    ...
2.10 区间遍历
2.10.1 全量迭代
__iter__ 直接返回 chain.from_iterable(self._lists),把逐块迭代交给 C 层的 chain。__reversed__ 先反转块顺序,再反转每块。两者都是 O(1)O(1) 构造迭代器,遍历本身 O(n)O(n)。使用 chain.from_iterable 比生成器表达式少一层 Python 帧。
2.10.2 按下标区间迭代
islice(start, stop, reverse) 先把 start、stop 正规化,然后用 _pos 定位起点和终点所在块号和块内下标。_islice 接收两个 (块号, 块内下标) 对,按块号关系分四种情况:
  1. 1.min_pos > max_pos:空区间;
  2. 2.min_pos == max_pos:同一块,直接块内切片;
  3. 3.next_pos == max_pos:相邻两块,拼接首块后缀与末块前缀;
  4. 4.跨多块:拼接首块后缀、中间整块、末块前缀。
正向用 chain 拼接,反向用 reversed(range(...)) 与 map(reversed, ...)。所有情况都用 map、chain、range,不在 Python 层写 for 循环。
2.10.3 按值区间迭代
irange(minimum, maximum, inclusive, reverse) 用两级二分把值区间转成下标区间,再复用 _islice。
2.11 其他细节
2.11.1 使用内置库
std 在热路径上尽量把循环交给内置函数:bisect、insort 负责二分与插入,chain.from_iterable 负责逐块迭代,map、starmap、reduce、iadd 负责切片拼接与建树。这些函数的循环体不在 Python 字节码层逐元素执行,CPython 下由 C 实现,单次调用的开销低于 Python 手写循环,但 PyPy 下可能无法应用 JIT 优化。
2.11.2 局部变量缓存
热路径方法开头都把属性缓存为局部变量,有的地方还把类方法也缓存下来。
Python
def _expand(self, pos):
    _load = self._load          # 缓存为局部变量,循环内少一次属性查找
    _lists = self._lists
    _index = self._index
    ...
Python
def __delitem__(self, index):
    ...
    _pos, _delete = self._pos, self._delete     # 循环内缓存方法查找
    for index in indices:
        pos, idx = _pos(index)
        _delete(pos, idx)
    ...
2.11.3 特判快路径
std 的快路径用一次比较取代 O(log⁡(n/L))O(\log(n/L)) 的树遍历或 O(L)O(L) 的块内移动。它们不改变渐进复杂度,只压低高频路径的常数。常见于:
  1. 1.首末块快速定位。 __getitem__ 与 pop 对下标落在首元素、末元素、首块、末块四种情况直接取块内元素,不调 _pos;_loc、_pos 也有对应的短路。
  2. 2.表尾快速添加。 查找值大于 _maxes[-1] 时,bisect_left、bisect_right 直接返回 _len,__contains__ 返回 False,add 把值追加到末块尾部。
2.12 复杂度小结
关于单次操作复杂度,文档与源码里同时存在两种说法:std 的 docstring 写的是“O(log(n)) -- approximate”;docs/performance-scale.rst 与 docs/paper.md 又说,把块长取到 n3\sqrt[3]{n} 时,均摊复杂度是 O(n3)O(\sqrt[3]{n})。这两种说法并不矛盾,因为它们描述的是不同的量。
单次插入的复杂度构成为:
  • 块间二分:O(log⁡(n/L))O(\log(n/L));
  • 块内二分:O(log⁡L)O(\log L);
  • 块内插入:O(L)O(L);
  • 索引更新:块数不变时单点更新 O(log⁡(n/L))O(\log(n/L));块数变化时整体重建 O(n/L)O(n/L),摊还后是 O(n/L2)O(n/L^2)。
合起来单次插入的摊还代价是
Tadd(n,L)=O(L+nL2+log⁡n)T_{\text{add}}(n, L) = O\left(L + \frac{n}{L^2} + \log n\right)

单次删除的推导和复杂度相同。bisect_left、bisect_right、__contains__ 不含块内移动,两级二分共 O(log⁡n)O(\log n);__getitem__ 只走索引,为 O(log⁡(n/L))O(\log(n/L))。
所以 O(log⁡n)O(\log n) 只描述两级二分定位这一步,忽略了块内移动与分裂;忽略对数项,当块长取 L=n3L=\sqrt[3]{n} 时,单次插入代价最低为 O(n3)O(\sqrt[3]{n}),但默认块长固定为 1000,没有随 nn 缩放,这个界不能直接套用。
此外,固定 L=1000L=1000 时,插入与删除的渐进复杂度是 O(n)O(n),但 n/L2n/L^2 一项超过 LL 时需要 n>L3=109n > L^3 = 10^9,在常见的数据规模下不会出现。因此固定块长在实践中够用,实际代价接近 O(L)O(L),不会表现出 O(n)O(n) 的增长。
3 PyRival 版 SortedList 分析
PyRival 是一个面向 Python 的算法竞赛模板库,其中有两个版本的 SortedList:cheran-senthil 编写的旧版(后文简称 pyr-cheran)与 bjorn-martinsson 主要编写的新版(后文简称 pyr-bjorn)。pyr-cheran 支持的类方法更全面,行数稍多;pyr-bjorn 行数更短,但基准中多数项目的用时高于 pyr-cheran。
3.1 树状数组构建索引
PyRival 最主要的改动是把 std 的线段树索引换成树状数组。线段树做下标换算、更新块长时固定遍历 log⁡B\log B 层;树状数组求前缀和只遍历 popcount(end)\text{popcount}(end) 次,平均 log⁡B/2\log B/2 步。但树状数组的树上二分要特判边界,所以实际运行效率不完全领先。
Python
def _fen_build(self):                   # O(B) 线性建树
    self._fen_tree[:] = self._list_lens
    for i in range(len(self._fen_tree)):
        if i | i + 1 < len(self._fen_tree):
            self._fen_tree[i | i + 1] += self._fen_tree[i]

def _fen_query(self, end):              # 前缀和,走 popcount(end) 步
    if self._rebuild:                   # 惰性重建
        self._fen_build()
    _fen_tree = self._fen_tree
    x = 0
    while end:
        x += _fen_tree[end - 1]
        end &= end - 1
    return x

def _fen_findkth(self, k):              # 树上二分,O(log B)
    ...
    _fen_tree = self._fen_tree
    idx = -1
    for d in reversed(range(len(_fen_tree).bit_length())):
        right_idx = idx + (1 << d)
        if right_idx < len(_fen_tree) and k >= _fen_tree[right_idx]:
            idx = right_idx
            k -= _fen_tree[idx]
    return idx + 1, k
pyr-bjorn 同样用树状数组,但把索引封装成单独的类,每次访问多一次属性查找和一层方法调用。
3.2 其余差异
pyr-cheran 用 _mins 存每块最小值作块间分隔值。bisect_left 落到的块,其首元素不小于 value,但 value 也可能属于前一块的尾部,定位后要多一步回退修正。pyr-bjorn 改用 macro 存下一块的首元素,二分直接给出目标块,不需要修正。
pyr-cheran 的索引重建仍沿用 std 的惰性重建,结构变化时只把 _rebuild 置为 True,下次查询再 _fen_build。pyr-bjorn 则每次块分裂就重建整棵 Fenwick 树,不做惰性重建。
两版都出于模板简洁性做了两项简化:不合并小块,块数没有上界;移除快路径,不做首尾特判。pyr-bjorn 的简化更多,代价是性能下降。
4 SortedList 重新实现与优化
基于以上分析,我们重写了一版 sortedlist_fast.py(以下简称 fast),其思路是沿用 std 的工程细节与 PyRival 的轻量结构,再针对 PyPy 做适配解释器的改动。
4.1 构建索引
重写 SortedList,先要确定用哪种树维护块长。下表是两种结构在三种操作上的相对用时,测试条件为 n=106n=10^6、块长 200 到 1000。
操作
树状数组 / 线段树
前缀和
0.36-0.39
单点更新
0.87-0.90
树上二分
1.21-1.29
前缀和一项,树状数组的用时只有线段树的 0.36-0.39 倍,bisect_left、bisect_right、index 都靠它把块坐标换算成全局下标;单点更新为 0.87-0.90 倍;树上二分则相反,为 1.21-1.29 倍,__getitem__、pop 依赖这一步,但首末块快路径也能覆盖一部分下标。综合来看树状数组占优,本实现用它构建块长索引。
本实现沿用 std 的惰性索引重建,只是把清空 _index 换成了置 dirty 标记:块数变化时打上重建标记,重建推迟到下一次读索引时进行。_fen_update 在 dirty 期间直接返回;_findkth、_fen_query 在 dirty 期间整体重建并清除标记。
4.2 块长取值
块长决定块内移动的长度,也决定块数与索引树的高度。不同版本对块长有不同选择:std 的 load 是 1000,pyr-cheran 的 _load 是 200,pyr-bjorn 的 block_size 是 700,块长达到 700 时对半分裂。不能简单套用已有数值,需要先测出 list 内元素移动的成本拐点。
对长度为 LL 的列表做两组测试:重复 b.insert(0, x) 直到长度变为 2L2L;重复 del b[0] 直到长度变为 0(用 del 而非 pop,省去返回值的开销)。下表是不同列表长度下的平均单次操作时间(ns/op):
列表长度
200
500
800
1000
1200
1500
2000
insert
43.7
87.4
121
143
205
265
348
del
27.2
29.1
36.3
41.4
329
760
1,350
insert 随列表长度近似线性增长;del 在 L≤1000L\le1000 时增长缓慢,但 L=1200L=1200 时从 41.4 跳到 329,拐点落在 1000 到 1200 之间。因此取 load=500,使块长上界 2L=10002L=1000 落在拐点以下。
4.3 小块合并
块被删除后会变短。std 在块长降到 L/2L/2 及以下时与邻块合并,块数因此稳定在 n/Ln/L 附近;合并会清空索引,等下一次读索引时重建。两版 PyRival 完全不合并,块数没有上界。本实现采取折中:不合并小块,块只在被删空时移除,平均块长低于 L/2L/2 时整体重建一次。
Python
def _delete(self, pos, idx):
    ...
    del block[idx]
    self.n -= 1
    if block:                            # 块非空:只更新块尾与索引
        maxes[pos] = block[-1]
        self._fen_update(pos, -1)
        return
    del lists[pos]                       # 块被删空:移除该块
    del maxes[pos]
    self.dirty = True                    # 块数变化,索引待重建
    if len(lists) > 1 and len(lists) * self.load > self.n << 1:
        self._compact()                  # 平均块长不足 L/2,整体重建
4.4 功能完善
除了提高性能以外,我们还想让 SortedList 对标 std::multiset 或者 pbds 的 tree,支持一些常用的类方法。
4.4.1 返回值
C++ 的 insert 返回指向新元素的迭代器,erase(iterator) 返回被删元素之后的迭代器。因此 fast 为 add、discard、remove 添加了返回值开关:默认关闭,行为与 std 一致;打开时返回插入或删除元素的原下标,discard 未命中时返回 -1,这样就不必再单独查一次下标。
Python
def add(self, value, index=False):
    ...
    res = self._fen_query(pos) + idx if index else None   # 开关打开时返回插入位置
    return res
4.4.2 区间迭代
C++ 遍历一段区间用两个迭代器,fast 则对下标区间迭代。islice(start, stop, reverse) 支持正向和逆向迭代区间,区间左闭右开、与 range 类似。islice 与 std 的结构类似,都依次迭代首块后缀、中间整块、末块前缀;区别在于 fast 用 Python 生成器逐块 yield from,std 用 chain 在 C 层拼接,因此 fast 的 islice 比 std 慢。
Python
def _iter_asc(self, start, stop):
    ...
    pos, idx = self._findkth(start)        # 定位起点所在块与块内下标
    block = lists[pos]
    yield from block[idx:]                 # 首块后缀
    ...
    while pos < len(lists):
        block = lists[pos]
        if end + len(block) >= stop:
            yield from block[:stop - end]  # 末块前缀
            return
        yield from block                   # 中间整块
        ...
4.4.3 区间删除
C++ 可以用两个迭代器逐个释放中间元素,fast 改用 del sl[a:b] 删除连续区间。step = 1 时,_del_range 依次裁掉首块尾部、裁掉尾块头部、移除中间整块,代价为 O(L+B)O(L + B),与区间长度无关;std 的切片删除只在删除量达到总量的 1/8 时整体重建,否则逐个下标删除。step != 1 时收集全部元素后整体重建。
Python
def __delitem__(self, index):
    ...
    if isinstance(index, slice):
        start, stop, step = index.indices(self.n)   # 正规化切片边界
        if step == 1:
            self._del_range(start, stop)            # 连续区间:按块裁剪
        else:
            values = self._collect()                # 非连续:收集全部元素
            del values[start:stop:step]
            self._rebuild(values)                   # 整体重建
        return
    pos, idx = self._pos(index)                     # 单下标:定位后删除
    self._delete(pos, idx)
4.5 其他细节
  • 手写二分内联,不调用 bisect。CPython 下内置函数更快;PyPy 下 JIT 把内联循环优化得比内置函数快 5% 左右。
  • 出于代码简洁性考虑,islice、iter 用生成器逐元素产出,没有沿用 std 的 chain 拼接。
  • 加入 __slots__,节省内存,也避免实例字典查找。
  • 树状数组二分时用 range(bits - 1, -1, -1) 而不是 reversed(range(bits)),后者多构造一层 reversed 迭代器,慢约 30%。
  • 参考 std 添加快路径。
4.6 复杂度小结
操作
复杂度
add、discard、remove
O(L+n/L2+log⁡n)O(L + n/L^2 + \log n)
__getitem__
O(log⁡(n/L))O(\log(n/L))
pop、__delitem__(单下标)
O(L+n/L2+log⁡n)O(L + n/L^2 + \log n)
__delitem__(切片、step = 1)
O(L+n/L)O(L + n/L)
bisect_left、bisect_right、index、count
O(log⁡n)O(\log n)
__contains__
O(log⁡n)O(\log n)
5 基准测试
5.1 测试标准
测试项目如下,前 14 项为单点操作,每次调用只处理一个元素;后 5 项为批量操作。
场景
操作
考察点
add
插入 10610^6 个随机值
随机位置插入,块内移动主导
add_asc
插入 10610^6 个严格递增的值
全部落在末尾,压测尾部追加快路径
add_dsc
插入 10610^6 个严格递减的值
全部落在块首,压测块首插入
discard
按值删除 10610^6 次,一半命中
命中走删除,未命中走定位
pop_mid
随机下标 pop 10610^6 次
随机位置删除
pop_front
pop(0) 10610^6 次
首块删除,压测块长收缩
pop_back
pop(-1) 10610^6 次
末块删除,走快路径
bisect_left
随机值查询 10610^6 次
两级二分 + 前缀和
bisect_right
随机值查询 10610^6 次
两级二分 + 前缀和
contains
查询 10610^6 次,50% 命中、50% 未命中
只做两级二分,不查索引
getitem
随机下标访问 10610^6 次
树上二分
index
按值查下标 10610^6 次,全部命中
两级二分 + 前缀和
mixed_write
10610^6 次随机 add、remove
纯写入混合负载
mixed_mix
10610^6 次随机 add、remove、bisect_left、getitem
读写混合负载
build
从 10610^6 个随机值构造
排序 + 切块
update
批量插入 10610^6 个新值
批量重建 vs 逐个插入
del_range
分 100 次删除随机一段,共删除 10610^6 个
按下标区间删除
islice
分 100 次迭代随机一段,共迭代 10610^6 个
按下标区间迭代
iter
完整遍历一次
迭代器开销
所有实验在 PyPy 3.10.13(v7.3.15 win64)下进行,每项 10 轮取平均;本机读数只用于版本间的相对比较。四份实现的接口并不一致:pyr-cheran 与 pyr-bjorn 都不提供 index、update、islice、del_range,基准用已有方法组合补全:index 为一次 bisect_left 加一次 getitem,update 逐个 add,islice 按下标逐个 getitem,del_range 从区间左端点起逐个 pop。
5.2 不同块长测试
先看 fast 版本在不同 load 下各项的用时(ns/op)。把每项按实际操作数还原成总时间后,L=300L=300 到 L=600L=600 的总时间为 3.20 到 3.25s,处于同一水平;更小的 load 抬高索引成本,更大的 load 抬高块内元素移动成本。因此取居中的 load=500,此时块长上界 2L=10002L=1000 恰好落在元素移动拐点之前。
场景
200
300
400
500
600
700
800
900
1000
add
390
414
427
439
454
477
495
498
512
add_asc
11.0
10.0
10.0
9.00
9.00
9.00
10.0
9.00
9.00
add_dsc
94.0
106
118
128
141
156
168
179
193
discard
203
200
201
199
202
203
205
208
206
pop_mid
263
252
249
242
244
246
243
246
247
pop_front
60.0
38.0
32.0
31.0
32.0
33.0
35.0
38.0
40.0
pop_back
12.0
11.0
11.0
11.0
11.0
11.0
10.0
10.0
10.0
bisect_left
288
282
287
279
276
273
274
276
276
bisect_right
292
290
281
278
277
275
276
279
273
contains
188
179
178
176
174
173
173
173
172
getitem
165
152
146
139
137
135
135
135
128
index
297
287
286
284
284
324
302
303
284
mixed_write
330
343
341
343
353
362
371
368
378
mixed_mix
344
339
334
339
345
342
345
343
350
build
134
133
133
132
136
133
133
133
132
update
164
163
163
164
165
164
164
162
163
del_range
1.00
0.640
0.500
0.430
0.390
0.340
0.320
0.290
0.280
islice
5.00
6.00
6.00
6.00
5.00
6.00
6.00
5.00
5.00
iter
7.00
5.00
5.00
4.00
4.00
4.00
4.00
4.00
4.00
总时间(s)
3.25
3.21
3.21
3.20
3.25
3.33
3.35
3.37
3.38
块长的影响仍然由成本模型 Lc1+log⁡(n/L)⋅c2L c_1 + \log(n/L)\cdot c_2 决定。插入类操作以块内元素移动为主,add 随 LL 增大而上升;索引类操作以树深为主,getitem、bisect_left、bisect_right 随 LL 增大而下降;删除类操作同时受两者影响。
5.3 不同版本测试
再看四份实现的对比。括号内为该项相对 std 的用时比值,小于 1 表示快于 std。
场景
std
fast
pyr-cheran
pyr-bjorn
add
484
449 (0.93)
475 (0.98)
592 (1.22)
add_asc
48.0
10.0 (0.20)
190 (3.92)
222 (4.58)
add_dsc
253
131 (0.52)
166 (0.66)
286 (1.13)
discard
188
198 (1.05)
246 (1.31)
546 (2.90)
pop_mid
265
245 (0.93)
485 (1.84)
295 (1.12)
pop_front
46.0
31.0 (0.68)
34.0 (0.74)
179 (3.89)
pop_back
11.0
11.0 (0.95)
19.0 (1.67)
149 (13.07)
bisect_left
368
279 (0.76)
293 (0.80)
331 (0.90)
bisect_right
373
283 (0.76)
295 (0.79)
344 (0.92)
contains
165
179 (1.08)
238 (1.44)
508 (3.08)
getitem
113
143 (1.27)
269 (2.38)
185 (1.64)
index
380
286 (0.75)
616 (1.62)
546 (1.44)
mixed_write
372
340 (0.91)
354 (0.95)
731 (1.96)
mixed_mix
369
341 (0.92)
430 (1.16)
580 (1.57)
build
134
134 (1.00)
138 (1.03)
134 (1.00)
update
166
167 (1.01)
478 (2.88)
594 (3.58)
del_range
139
0.500 (0.00359)
331 (2.38)
182 (1.31)
islice
2.00
4.00 (2.41)
207 (121.67)
122 (71.71)
iter
1.00
4.00 (3.00)
4.00 (2.92)
4.00 (2.92)
综合来看,fast 在多数插入、删除与定位类项目上快于 std,落后较多的是 getitem、islice、iter,discard 与 contains 略慢于 std。pyr-cheran 仅有少数几项快于 std,大部分项目都慢于 std。islice 与 iter 的差距来自接口实现:pyr 两版没有 islice,只能逐个 getitem;fast 的 __iter__ 与 islice 用生成器逐元素产出,std 用 chain.from_iterable 在 C 层拼接。del_range 的差距来自删除方式:fast 走块切片,std 在这个数据规模下退回逐个下标删除,pyr 两版从段首逐个 pop。
四份实现的差异可以归纳为增删、索引与批量处理三方面:增删类差距来自块的维护方式,索引类差距来自树形索引结构,批量处理差距来自接口是否提供对应方法。
5.4 在线评测
我们选用了几道模板题目作为 OJ 环境的测试。
洛谷 P6136 【模板】普通平衡树(数据加强版)。提交 用时 16.14s,目前是该题最快的提交,快于直接使用 pyr-cheran 的提交。
Library Checker Ordered Set。提交 用时 671ms,慢于使用分块 SortedSet 的提交(580ms),快于 AVL 的提交(710ms)和离线离散化树状数组的提交(711ms)。
差距主要来自数据结构本身的设计。SortedList 维护的是有序多重集合,支持的功能是有序集合的超集,因此在一些针对有序集合的特化操作上会产生额外开销:添加元素时,需要先判断是否存在、再重新插入;查找前驱后继时,需要先查找下标、再重新访问。此外,根据题目的数据生成器,测试数据均为随机,无法反映出在 pop_front 等极端场景下的表现。更快的 SortedSet 在分块后不维护块长索引,定位时线性扫描所有块,单次操作代价随块数增长,只适合块数较少的场景。
CodeForces 896C,也是珂朵莉树模板题。提交 用时 1984ms。这里继承 SortedList 重写了一版 SortedSet,依靠常数优化和评测机波动极限通过,是 Python 提交中少数通过该题的代码之一。
6 总结
SortedList 的设计思路可以总结为:用连续内存的 C 层移动替代解释器层的指针追踪;用分块把单次插入要移动的元素数从 nn 降到块长 LL;用树状数组把下标换算的代价降到 O(log⁡B)O(\log B);用惰性重建把结构维护的成本摊还到每次操作。fast 基于 std 的工程细节与 pyr-cheran 的轻量结构实现,并针对 PyPy 做了适配优化;代码量有所增加,但多数项目用时低于 std。
本文的代码、测试、分析与写作由 DeepSeek-V4.1-Flash 辅助完成。
相关阅读
评论 (4)