已排序整数列表查找数值变化索引的高效实现方案咨询

问题背景 #

假设我们有如下已排序整数列表：

data = [1] * 3 + [4] * 5 + [5] * 2 + [9] * 3
# 展开后为 [1, 1, 1, 4, 4, 4, 4, 4, 5, 5, 9, 9, 9]

需要找出列表中数值发生变化的分界索引，预期输出为：

[3, 8, 10, 13]

现有两种实现方案：

1. 基于itertools.groupby的实现：

from itertools import groupby

cursor = 0
result = []
for key, group in groupby(data):
    cursor += sum(1 for _ in group)
    result.append(cursor)
print(result)

输出和预期一致，时间复杂度为O(n)。
2. 基于bisect.bisect_left的实现：

from bisect import bisect_left

cursor = 0
result = []
while cursor < len(data):
    cursor = bisect_left(data, data[cursor] + 1, cursor, len(data))
    result.append(cursor)
print(result)

输出和预期一致，时间复杂度为O(k*log n)，其中k为列表中不同元素的数量，也可以替换为指数搜索作为变体。

提问：是否有更快、性能更优的实现方式？ #

解答 #

是否有更优实现要看使用场景和可依赖的工具，分情况说明：

纯Python标准库场景 #

• 先明确两种现有方案的适用边界：
  • 如果列表中不同元素的数量k非常小（属于大量重复元素的有序列表），bisect方案的O(k logn)会比遍历类的O(n)方案快得多，比如百万级元素只有十几个不同值时，bisect方案的实际运行速度是遍历方案的几十上百倍。
  • 如果列表中元素重复率很低，k接近总长度n，bisect方案会退化为O(n logn)，此时遍历类方案性能更好。
• 可以优化遍历类方案的实现，用Python内置的C实现迭代器减少Python层面的循环开销，比原始groupby写法性能高30%以上：

# 相邻元素比较法，底层用zip和列表推导，核心逻辑都是C实现
result = [i + 1 for i, (a, b) in enumerate(zip(data, data[1:])) if a != b] + [len(data)]

可依赖第三方数值库场景 #

如果可以用numpy，性能会有数量级的提升，直接用向量化运算处理，完全避免Python层面的循环：

import numpy as np

arr = np.array(data)
# 求相邻元素的差异，找到差异不为0的位置后加1得到分界点
result = np.where(np.diff(arr) != 0)[0] + 1
# 补上最后一个分界点（列表总长度）
result = np.append(result, len(arr)).tolist()

这种实现对于十万级以上的列表，性能比纯Python实现快几十到上百倍，非常适合大规模数据处理。

特殊场景优化 #

如果你的列表元素是值域很小的整数，还可以用计数统计+前缀和的方式计算分界点，性能会比bisect方案还要高，但仅适合元素取值范围有限的场景。 #

内容的提问来源于stack exchange，提问作者Dani Mesejo