如何在Pandas中实现超另一列累计值后重置调整的分组累计求和

核心逻辑说明 #

你需要的是依赖前一步结果的自定义运行总计：每一步将当前a的值累加到前一步的调整后总计中，若结果超过当前b的运行总计，则将其调整为b的运行总计（即减去超出部分）。由于这种逻辑存在状态依赖，无法直接用Pandas原生的cumsum实现，需用循环或JIT加速的方式处理。

1. 基础案例实现 #

先针对你提供的初始DataFrame实现核心逻辑：

步骤1：定义调整函数 #

def adjusted_running_total(a_series, b_rt_series):
    adjusted = []
    prev_total = 0
    for a_val, b_rt_val in zip(a_series, b_rt_series):
        current_total = prev_total + a_val
        # 超出时重置为b的运行总计
        if current_total > b_rt_val:
            current_total = b_rt_val
        adjusted.append(current_total)
        prev_total = current_total
    return adjusted

步骤2：应用到DataFrame #

import pandas as pd

data = {'date': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04','2023-01-05','2023-01-06'],
        'a': [10, 10, 15, 15,15,15],
       'b': [0,5,60,5,0,80]
       }
df = pd.DataFrame.from_dict(data)

# 计算b的运行总计
df['b rt'] = df['b'].cumsum()
# 计算调整后的a运行总计
df['adjusted a rt'] = adjusted_running_total(df['a'], df['b rt'])

运行后得到的结果：

2. 按product/location分组实现 #

如果需要对每个product和location的组合单独计算，只需在分组后应用上述逻辑：

示例输入数据 #

data = {
    'product': ['A', 'A', 'A', 'B', 'B', 'B'],
    'location': ['X', 'X', 'X', 'Y', 'Y', 'Y'],
    'date': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03'],
    'a': [10, 10, 15, 5, 5, 10],
    'b': [0,5,60,2,3,15]
}
df = pd.DataFrame(data)

分组处理代码 #

def process_group(group):
    # 计算当前分组的b运行总计
    group['b rt'] = group['b'].cumsum()
    # 计算当前分组的调整后a运行总计
    group['adjusted a rt'] = adjusted_running_total(group['a'], group['b rt'])
    return group

# 按product和location分组处理，保持原DataFrame结构
df = df.groupby(['product', 'location'], group_keys=False).apply(process_group)

运行后每个分组会独立计算逻辑，结果示例：

3. 大数据量优化（可选） #

如果你的DataFrame行数较多（十万级以上），纯Python循环会较慢，可使用numba进行JIT加速：

from numba import jit
import numpy as np

@jit(nopython=True)
def fast_adjusted_rt(a_array, b_rt_array):
    n = len(a_array)
    adjusted = np.zeros(n, dtype=np.int64)
    prev_total = 0
    for i in range(n):
        current_total = prev_total + a_array[i]
        if current_total > b_rt_array[i]:
            current_total = b_rt_array[i]
        adjusted[i] = current_total
        prev_total = current_total
    return adjusted

# 修改分组处理函数使用加速版本
def process_group_fast(group):
    group['b rt'] = group['b'].cumsum()
    group['adjusted a rt'] = fast_adjusted_rt(group['a'].values, group['b rt'].values)
    return group

df = df.groupby(['product', 'location'], group_keys=False).apply(process_group_fast)

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