在Pandas DataFrame中按Streak_Count指定数量求和MON TOTAL列

嘿，这个需求我刚好处理过，用Pandas可以轻松实现任意数量的连续求和！咱们先理清楚逻辑：当STREAK_COUNT大于0时，需要计算当前行及往前STREAK_COUNT-1行的TOTAL之和；等于0时结果为空。

方法一：使用apply逐行处理（直观易懂） #

先把你的示例数据转换成DataFrame方便测试：

import pandas as pd

data = {
    'MON': ['1/2/1992', '2/3/1992', '3/2/1992', '4/1/1992', '5/1/1992', '7/1/1992', '8/3/1992', '10/1/1992', '11/2/1993', '1/4/1993', '2/1/1993', '3/1/1993', '4/1/1993'],
    'TOTAL': [1.123077, -1.296718, -6.355612, 5.634692, 4.180605, -0.101016, -0.706125, 0.368579, 3.822277, 2.233359, 15.219644, -2.647693, 1.599094],
    'STREAK_COUNT': [1, 0, 2, 0, 2, 0, 2, 0, 0, 0, 4, 1, 1],
    'WANTED RESULT': [1.123077, None, -7.65233, None, 9.815297, None, -0.807141, None, None, None, 21.643859, -2.647693, 1.599094]
}

df = pd.DataFrame(data)

然后定义一个自定义函数，逐行计算求和：

def calculate_streak_sum(row):
    streak = row['STREAK_COUNT']
    if streak == 0:
        return None
    # 计算求和的起始索引，确保不小于0
    start_idx = max(row.name - streak + 1, 0)
    # 取对应区间的TOTAL值求和
    return df.loc[start_idx : row.name, 'TOTAL'].sum()

# 应用函数生成结果列
df['CALCULATED_RESULT'] = df.apply(calculate_streak_sum, axis=1)

方法二：使用前缀和（高效处理大数据量） #

如果你的DataFrame行数很多，apply的效率会偏低，这时候可以用前缀和的方法，时间复杂度是O(n)，比方法一更高效：

# 计算TOTAL列的前缀和
prefix_sum = df['TOTAL'].cumsum()

# 利用前缀和快速计算区间和
df['CALCULATED_RESULT'] = df.apply(lambda row: 
    # 当streak>0时，用当前前缀和减去起始位置的前缀和（如果起始位置>0）
    prefix_sum[row.name] - (prefix_sum[row.name - row['STREAK_COUNT']] if (row.name - row['STREAK_COUNT']) >= 0 else 0)
    if row['STREAK_COUNT'] > 0 else None, axis=1)

验证结果 #

运行后你会发现CALCULATED_RESULT列和你给出的WANTED RESULT完全一致：

• 当STREAK_COUNT=4时，计算的是连续4个TOTAL值的和：0.368579 + 3.822277 + 2.233359 + 15.219644 = 21.643859
• 当STREAK_COUNT=2时，计算的是当前行和上一行的TOTAL之和，比如第三行的-1.296718 + (-6.355612) = -7.65233

两种方法都能完美实现你要的任意数量求和需求，根据你的数据规模选择就行~

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