pandas时序数据中替代逐行迭代计算行累计值的高效实现方法

我有如下代码，用于根据预算计算可购买的商品数量，当前实现用了逐行迭代逻辑，我希望将其优化为更符合pandas风格的高效写法，避免逐行迭代、避免重复计算截至当前行的累计值。
原代码如下：

import math
import pandas as pd

data = [['2021-01-02', 5.5], ['2021-02-02', 10.5], ['2021-03-02', 15.0], ['2021-04-02', 20.0]]
df = pd.DataFrame(data, columns=['Date', 'Current_Price'])

df.Date = pd.to_datetime(df.Date)
mn = df.Date.min()
mx = df.Date.max()
dr = pd.date_range(mn - pd.tseries.offsets.MonthBegin(), mx + pd.tseries.offsets.MonthEnd(), name="Date")
df = df.set_index("Date").reindex(dr).reset_index()
df['Current_Price'] = df.groupby(
    pd.Grouper(key='Date', freq='1M'))['Current_Price'].ffill().bfill()

# 新增Day列标识当前日期是当月的第几天
df['Day'] = pd.to_datetime(df['Date']).dt.day

# 初始化目标列
df[['Deposit', 'Withdrawal', 'Num_of_Products_Bought']] = 0.0

# 每月月初存入100美元
df.loc[df['Day'] == 1, 'Deposit'] = 100.0

原迭代逻辑如下：

for index, row in df.iterrows():
    if df.loc[index, 'Day'] == 1:
        # 购买商品数量 = (截至当前累计存款 - 截至当前累计支出) / 当前商品价格
        df.loc[index, 'Num_of_Products_Bought'] = math.floor(
            (sum(df.iloc[0:(index + 1)]['Deposit'])
             - sum(df.iloc[0:(index + 1)]['Withdrawal']))
            / df.loc[index, 'Current_Price'])
        # 记录对应日期购买商品的总支出
        df.loc[index, 'Withdrawal'] = df.loc[index, 'Num_of_Products_Bought'] * df.loc[index, 'Current_Price']

你提到的cumsum不能直接用是因为每一期的支出会影响后续的可用余额，属于带状态的迭代计算，我们可以通过只处理有交易的日期、搭配滚动计算来实现提速，完全避免全表逐行遍历：

纯pandas风格实现 #

# 筛选出每月1号的交易行，其他行目标列默认是0无需处理
trade_df = df[df['Day'] == 1].reset_index(drop=False).copy()
# 累计存款可以直接用cumsum预计算
trade_df['cum_deposit'] = trade_df['Deposit'].cumsum()

# 自定义滚存计算函数，处理每一期的购买数量和支出
def calc_trade(window):
    # 当前窗口最后一行就是当前待计算的行
    current_row = window.iloc[-1]
    if len(window) == 1:
        prev_withdrawal_sum = 0
    else:
        # 之前所有行的支出总和
        prev_withdrawal_sum = window.iloc[:-1]['Withdrawal'].sum()
    available = current_row['cum_deposit'] - prev_withdrawal_sum
    num = math.floor(available / current_row['Current_Price'])
    withdraw = num * current_row['Current_Price']
    return pd.Series([num, withdraw], index=['Num_of_Products_Bought', 'Withdrawal'])

# 对交易行用expanding窗口计算
trade_df[['Num_of_Products_Bought', 'Withdrawal']] = trade_df.expanding().apply(calc_trade)[['Num_of_Products_Bought', 'Withdrawal']]

# 把计算结果写回原df
df.loc[trade_df['index'], ['Num_of_Products_Bought', 'Withdrawal']] = trade_df[['Num_of_Products_Bought', 'Withdrawal']].values

超高性能精简实现 #

如果数据量特别大，还可以进一步用变量滚存计算交易行，再合并回df，性能会更好：

trade_df = df[df['Day'] == 1].reset_index(drop=False).copy()
cum_deposit = 0
cum_withdraw = 0
nums = []
withdraws = []
for _, row in trade_df.iterrows():
    cum_deposit += row['Deposit']
    available = cum_deposit - cum_withdraw
    num = math.floor(available / row['Current_Price'])
    withdraw = num * row['Current_Price']
    nums.append(num)
    withdraws.append(withdraw)
    cum_withdraw += withdraw
trade_df['Num_of_Products_Bought'] = nums
trade_df['Withdrawal'] = withdraws
df.loc[trade_df['index'], ['Num_of_Products_Bought', 'Withdrawal']] = trade_df[['Num_of_Products_Bought', 'Withdrawal']].values

两种写法的计算结果和原迭代逻辑完全一致，但性能远高于全表逐行遍历：

• 第一种是纯pandas风格实现，不需要自己维护状态变量
• 第二种仅遍历有交易的行（总共有N个月就只遍历N次，远低于原逻辑的总天数次遍历）

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