Pandas：如何计算当前行与最近满足条件行的行数间隔？

需求说明 #

需要在DataFrame中生成列b，其值为当前行距离列a最后一次为True位置的间隔数，示例数据如下：

a       b
0    True    0
1    False   1
2    True    0
3    False   1
4    False   2
5    False   3

当前使用循环实现的代码运行速度极慢，代码如下：

import numpy as np
import pandas as pd

# 假设cond是生成列a的判断条件
a = np.where(cond)[-1]
b = np.array([], dtype=np.int64)
s = 0
for i in range(0, len(data)):
    if i in a:
        b = np.append(b, 0)
        s = 0
    else:
        b = np.append(b, s)
    s += 1
data['b'] = pd.Series(b).fillna(method='ffill').fillna(-1)

高效非循环实现方案 #

方法一：分组累计计数 #

通过cumsum()标记分组，再在每个分组内从0开始计数：

import pandas as pd

# 生成分组ID：每遇到a为True时分组ID自增
data['group'] = data['a'].cumsum()
# 每个分组内从0开始累计计数
data['b'] = data.groupby('group').cumcount()
# 移除辅助分组列（可选）
data.drop('group', axis=1, inplace=True)

方法二：计算索引差值 #

通过填充最近的True位置，再计算当前索引与该位置的差值：

import pandas as pd

# 标记所有True的索引位置，并向前填充空值
data['last_true_idx'] = data.index.where(data['a']).ffill()
# 计算当前行与最近True位置的间隔
data['b'] = data.index - data['last_true_idx']
# 移除辅助列（可选）
data.drop('last_true_idx', axis=1, inplace=True)

优势说明 #

两种方法均采用pandas内置的向量化操作，完全避免循环。在处理大数据量时，速度比原循环实现提升显著，尤其当DataFrame行数超过10万时，性能差距会达到数倍甚至数十倍。

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