如何基于Pandas DataFrame列值变化高效计算COUNT与SUM列？

需求说明 #

需要对大体积DataFrame按SIGNAL列的变化计算两个统计值：

1. 每次SIGNAL从0切换为1时，统计本次切换行和上一次SIGNAL=1的行之间的总长度作为COUNT
2. COUNT不为0时，计算当前行与上一次SIGNAL=1行的VALUE之和作为SUM

实现逻辑 #

采用全向量化操作，完全避免遍历，性能适配百万级以上的大数据量场景：

1. 先将VALUE列转换为数值类型，避免字符串类型导致求和异常
2. 提取所有SIGNAL=1的行的索引和对应VALUE值
3. 计算相邻SIGNAL=1行的索引差、VALUE和，回填到对应行
4. 其余位置COUNT、SUM统一填充为0

完整可运行代码 #

import pandas as pd
import numpy as np

# 初始化数据
data = [['1', 0], ['2', 1], ['3', 0], ['4', 1], ['5', 0], ['6', 0], ['7', 0], ['8', 1], ['9', 0], ['10', 0], ['11', 0]]
df = pd.DataFrame(data, columns = ['VALUE', 'SIGNAL'])
# VALUE列转数值，原初始化数据为字符串会影响求和
df['VALUE'] = pd.to_numeric(df['VALUE'])

# 初始化统计列为0
df['COUNT'] = 0
df['SUM'] = 0

# 筛选所有SIGNAL=1的位置和对应VALUE
signal_1_mask = df['SIGNAL'] == 1
signal_1_indices = df.index[signal_1_mask]
signal_1_values = df.loc[signal_1_mask, 'VALUE'].values

# 至少有2个SIGNAL=1时才需要计算统计值
if len(signal_1_indices) >= 2:
    # 计算COUNT：当前1的索引 - 上一个1的索引 +1，匹配需求示例逻辑
    counts = signal_1_indices[1:] - signal_1_indices[:-1] + 1
    # 计算SUM：当前VALUE + 上一个1的VALUE
    sums = signal_1_values[1:] + signal_1_values[:-1]
    # 回填到对应行
    df.loc[signal_1_indices[1:], 'COUNT'] = counts
    df.loc[signal_1_indices[1:], 'SUM'] = sums

print(df)

运行结果 #

VALUE  SIGNAL  COUNT  SUM
0       1       0      0    0
1       2       1      0    0
2       3       0      0    0
3       4       1      3    6
4       5       0      0    0
5       6       0      0    0
6       7       0      0    0
7       8       1      5   12
8       9       0      0    0
9      10       0      0    0
10     11       0      0    0

问题修复说明 #

你之前版本的代码错误原因是采用全局信号切换次数累计的逻辑，没有仅计算相邻SIGNAL=1行的间隔，当SIGNAL频繁切换时就会出现计数错误。本方案直接针对SIGNAL=1的行计算间隔，不受中间切换次数影响，可适配所有场景。

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