如何在Pandas中按类别计算带max的累加和（无循环）

问题背景 #

需要按category分组计算sum字段，递推规则为：
sumₙ = max(0, diffₙ + sumₙ₋₁ - factor)，其中factor=2，且hour=0时sum初始值为0。原始数据和预期输出如下：

原始DataFrame

预期输出

由于数据量较大，需避免iterrows或全局循环，而无max(0, ...)约束时可用groupby.cumsum，但截断逻辑导致该方法失效，需针对性解决方案。

解决方案1：组内循环（适用于中等数据量） #

利用groupby.apply在每个分组内执行局部循环，比全局循环效率更高：

import pandas as pd

# 构造原始数据
df = pd.DataFrame({
    'category': ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b'],
    'hour': [0, 1, 2, 3, 0, 1, 2, 3],
    'diff': [0, 4, 3, 1, 0, 1, -5, 4],
    'sum': [0, pd.NA, pd.NA, pd.NA, 0, pd.NA, pd.NA, pd.NA]
})

factor = 2

def calculate_group_sum(group):
    sum_vals = group['sum'].copy()
    # 从第2行开始递推计算
    for idx in range(1, len(group)):
        prev_sum = sum_vals.iloc[idx-1]
        current_diff = group['diff'].iloc[idx]
        sum_vals.iloc[idx] = max(0, current_diff + prev_sum - factor)
    return sum_vals

# 按category分组计算
df['sum'] = df.groupby('category', group_keys=False).apply(calculate_group_sum)

解决方案2：Numba加速（适用于大数据量） #

当数据量极大时，纯Python循环仍有性能瓶颈，可借助Numba将循环编译为机器码，大幅提升速度：

import pandas as pd
import numpy as np
from numba import jit

# 构造原始数据
df = pd.DataFrame({
    'category': ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b'],
    'hour': [0, 1, 2, 3, 0, 1, 2, 3],
    'diff': [0, 4, 3, 1, 0, 1, -5, 4],
    'sum': [0, pd.NA, pd.NA, pd.NA, 0, pd.NA, pd.NA, pd.NA]
})

@jit(nopython=True)
def numba_sum_calculator(diff_array, initial_val):
    n = len(diff_array)
    sum_array = np.zeros(n, dtype=np.float64)
    sum_array[0] = initial_val
    for i in range(1, n):
        sum_array[i] = max(0.0, diff_array[i] + sum_array[i-1] - 2.0)
    return sum_array

# 分组应用Numba加速函数
df['sum'] = df.groupby('category').apply(
    lambda group: numba_sum_calculator(group['diff'].values, group['sum'].iloc[0])
).explode().values

关键说明 #

• 无法直接使用cumsum的原因：当sum被截断为0后，后续计算依赖的是截断后的值，而非未截断的累积和，因此必须通过递推实现。
• Numba版本的优势：nopython=True模式会将函数编译为机器码，执行效率接近原生C代码，适合百万级以上数据量的场景。

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