如何在Pandas中向量化计算填充available_wheat列？规避iterrows

问题回顾 #

需要基于初始的available_wheat值，按usage rate (%)的比例逐行递减计算后续的available_wheat，同时算出consumption，且要替代效率低下的iterrows循环，用向量化方法处理大数据量。

向量化实现代码 #

直接利用Pandas的累积乘积（cumprod）实现向量化计算，完全避免循环，效率大幅提升：

import pandas as pd

# 构造示例DataFrame（实际使用时替换为你的数据）
df = pd.DataFrame({
    'County': ['A']*12,
    'date': ['1/2/2021', '1/3/2021', '1/4/2021', '1/5/2021', '1/6/2021', '1/7/2021',
             '1/8/2021', '1/9/2021', '1/10/2021', '1/11/2021', '1/12/2021', '1/13/2021'],
    'available_wheat': [100.00] + [None]*11,
    'usage rate (%)': [3,3,2,5,1,2,5,6,7,8,1,2],
    'consumption': [None]*12
})

# 1. 计算每一期的留存率（1 - 使用率）
retention = 1 - df['usage rate (%)'] / 100

# 2. 填充available_wheat：初始值 × 前序所有留存率的累积乘积
# shift(1)让每一期使用上一期的留存率，fillna(1)保证第一期用初始值
df['available_wheat'] = df.loc[0, 'available_wheat'] * retention.shift(1).fillna(1).cumprod()

# 3. 计算当期消耗：当期可用量 × 当期使用率
df['consumption'] = df['available_wheat'] * (df['usage rate (%)'] / 100)

# 保留两位小数，匹配预期输出
df = df.round(2)

多分组场景适配 #

如果你的数据包含多个County（需要按分组独立计算），可以结合groupby和自定义函数实现：

def calculate_wheat_group(group):
    # 获取当前分组的初始可用小麦量
    initial_val = group.loc[group['available_wheat'].first_valid_index(), 'available_wheat']
    retention = 1 - group['usage rate (%)'] / 100
    # 按分组计算可用量和消耗量
    group['available_wheat'] = initial_val * retention.shift(1).fillna(1).cumprod()
    group['consumption'] = group['available_wheat'] * (group['usage rate (%)'] / 100)
    return group.round(2)

# 按County分组计算
df = df.groupby('County').apply(calculate_wheat_group)

效果说明 #

运行后得到的结果完全匹配你给出的预期输出，且向量化操作的时间复杂度远低于循环，在大数据量下（比如百万级行），速度会比iterrows快几十甚至上百倍。

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