如何在Pandas DataFrame分组内统计非零值出现后的行数？

现有如下DataFrame：

需要按ID分组，添加一列num_rows_since_nonzero，统计当前行距离最近一次非零值出现的行数：

• 非零值所在行的计数为0
• 非零值之后的零值，依次递增计数
• 分组开头的零值（未遇到任何非零值前）计数为0

预期结果：

可以通过pandas的分组和矢量化操作实现，无需循环，效率较高：

import pandas as pd

# 构造原始数据
df = pd.DataFrame({
    'ID': ['A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'B'],
    'value': [0, 1, 0, 0, 0, 0, 2, 0, 4, 0]
})

# 计算目标列
df['num_rows_since_nonzero'] = df.groupby('ID').transform(
    lambda x: x['value'].ne(0).cumsum().pipe(lambda s: s.groupby(s).cumcount()).where(x['value'].eq(0), 0)
)

print(df)

1. 标记非零位置：x['value'].ne(0)生成布尔序列，非零值对应True，零值对应False
2. 划分连续块：cumsum()对布尔序列累加，每遇到一个非零值，累加值就会+1，这样每组内被划分为多个连续块（每个块从一个非零值开始，到下一个非零值前结束）
3. 块内计数：pipe(lambda s: s.groupby(s).cumcount())对每个块内的行从0开始计数，非零值所在行计数为0，后续零值依次递增
4. 重置非零值计数：where(x['value'].eq(0), 0)确保非零值所在行的计数保持为0，符合需求

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