Pandas计算序列相同值首次出现间隔的高效实现方法

你当前使用的两层嵌套循环时间复杂度为O(n²)，数据量超过万级时性能会明显下降，有两种更高效的实现方案，时间复杂度最低可达O(n)：

方案1：倒序遍历+哈希表（易理解，性能稳定） #

思路：从后往前遍历A列，用字典存储每个值最后一次出现的索引，遍历到当前元素时如果字典中存在该值，直接计算索引差值存入B列，最后更新字典中该值的索引即可，全程仅需遍历一次数据。

import pandas as pd
import numpy as np

df = pd.DataFrame(columns=list("AB"))
df["A"] = [4,2,4,4,2,5,2,6,1,6,4,5,9,3,7,3,3]

# 初始化B列为空值
df['B'] = np.nan
last_occur = dict()

# 从最后一行往前遍历
for idx in reversed(df.index):
    val = df.loc[idx, 'A']
    if val in last_occur:
        df.loc[idx, 'B'] = last_occur[val] - idx
    # 更新当前值的最后出现索引
    last_occur[val] = idx

方案2：Pandas向量化操作（代码最简洁，大数据量性能最优） #

思路：利用Pandas的分组偏移特性，对A列相同值的组内索引做后一位偏移，直接和当前索引做差得到结果，全程无Python级循环，性能远高于自定义遍历逻辑。

import pandas as pd

df = pd.DataFrame({'A': [4,2,4,4,2,5,2,6,1,6,4,5,9,3,7,3,3]})
# 按A列分组，取组内下一个元素的索引减去当前索引得到间隔
df['B'] = df.groupby('A').apply(lambda x: x.index.to_series().shift(-1) - x.index).reset_index(level=0, drop=True)

两种方案的适用场景：

• 倒序遍历方案：适合需要自定义特殊计算逻辑的场景，代码逻辑好调试，性能足够应对百万级以内的数据
• 向量化方案：适合纯数值计算场景，代码量少，千万级数据也能快速处理

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