求助：高效实现带多条件层级的DataFrame批量计算方法

问题背景 #

现有一个按item_ID分组、日期排序的pandas DataFrame，结构如下：

import pandas as pd

df = pd.DataFrame(
 {
   "date": ["2020-01-01", "2020-02-01", "2020-03-01", "2020-01-01", "2020-02-01"],
   "item_ID": ["a", "a", "a", "b", "b"],
   "quantity" : [1, 2, 3, 5, 2],
   "type": ["y", "y", "n", "y", "n"],
 }
)

数据展示：

date    | item_ID | quantity | type |
2020-01-01  |    a    |    1     |  y   |
2020-02-01  |    a    |    2     |  y   |
2020-03-01  |    a    |    3     |  n   |
2020-01-01  |    b    |    5     |  y   |
2020-02-01  |    b    |    2     |  n   |

处理规则 #

• 当遇到type为"n"的行时，从该item_ID之前日期的type为"y"的行的quantity中扣减该行的quantity，且quantity不能为负。
• 若扣减后某行quantity为负，则将该行quantity设为0，剩余扣减值继续向前找更早的type为"y"的行扣减，直到"n"类行的扣减值耗尽。
• 最终移除所有type为"n"的行。

目标输出 #

date    | item_ID | quantity | type |
2020-01-01  |    a    |    0     |  y   |
2020-02-01  |    a    |    0     |  y   |
2020-01-01  |    b    |    3     |  y   |

当前采用循环遍历每个item_ID的方法处理大数据集时效率极低，需寻求更高效的实现方案。

高效实现方案 #

利用pandas的分组、累积计算和向量化操作替代循环，大幅提升处理效率：

import pandas as pd

def process_group(group):
    # 分离y和n类型的行
    y_rows = group[group['type'] == 'y'].copy()
    n_rows = group[group['type'] == 'n'].copy()
    
    if n_rows.empty:
        return y_rows
    
    # 计算y行的累积数量（按日期从早到晚）
    y_rows['cum_quantity'] = y_rows['quantity'].cumsum()
    # 计算当前组所有n行的总扣减值
    total_deduct = n_rows['quantity'].sum()
    
    # 计算扣减后的剩余累积量，小于0则取0
    y_rows['remaining'] = y_rows['cum_quantity'].sub(total_deduct).clip(lower=0)
    # 通过相邻剩余量的差值，得到每行实际剩余的quantity
    y_rows['quantity'] = y_rows['remaining'].sub(y_rows['remaining'].shift(fill_value=0))
    
    # 保留目标列并返回
    return y_rows[['date', 'item_ID', 'quantity', 'type']]

# 按item_ID分组处理，合并结果后重置索引
result = df.groupby('item_ID', group_keys=False).apply(process_group).reset_index(drop=True)
print(result)

代码说明 #

1. 分组隔离：通过groupby('item_ID')确保每个商品的扣减逻辑独立处理，不会互相干扰。
2. 类型分离：将组内的"y"和"n"行拆分，简化后续计算逻辑。
3. 累积计算优化：
   • 对"y"行计算累积数量，快速定位扣减后的剩余边界。
   • 直接计算所有"n"行的总扣减值，避免逐行处理的冗余操作。
4. 向量化扣减：利用sub、clip、shift等向量化方法批量计算剩余数量，完全替代Python循环的逐行判断，大幅降低计算开销。
5. 结果整合：将所有组处理后的"y"行合并，得到最终符合要求的数据集。

该方案完全基于pandas的内置优化操作，在大数据集上的处理效率会比循环方法提升数倍甚至数十倍。

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