优化Pandas DataFrame交易资金计算：寻求非逐行循环方案

需求说明 #

当前通过逐行循环计算每笔交易后的资金数据，希望改用Pandas向量化操作优化，核心逻辑：

• TRADE=1：买入，用当前CAPITAL除以<Open>价得到POSITION，GAIN设为0，GAIN_C继承上一行值
• TRADE=-1：卖出，计算<Open>价与上一笔买入价的DIFF，用DIFF*上一行POSITION得到GAIN，更新GAIN_C和CAPITAL，POSITION重置为0

原循环实现代码 #

for i in range(1, len(df_trades)):
    if df_trades.loc[i, 'TRADE'] != -1: # BUY order: set POSITION for buy order
        df_trades.loc[i, 'POSITION'] = df_trades.loc[i,'CAPITAL']//df_trades.loc[i,'<Open>']
        df_trades.loc[i,'GAIN'] = 0
        df_trades.loc[i,'GAIN_C'] = df_trades.loc[i -1,'GAIN_C']
    if df_trades.loc[i, 'TRADE'] == -1: # SELL order: recalculate Capital, gain, gain_c
        df_trades.loc[i, 'POSITION'] = 0
        df_trades.loc[i,'DIFF'] = df_trades.loc[i,'<Open>'] - df_trades.loc[i-1,'<Open>']            
        df_trades.loc[i,'GAIN'] = df_trades.loc[i,'DIFF'] * df_trades.loc[i-1,'POSITION']
        df_trades.loc[i,'GAIN_C'] = df_trades.loc[i-1,'GAIN_C'] + df_trades.loc[i,'GAIN']
        df_trades.loc[i,'CAPITAL'] = df_trades.loc[i-1,'CAPITAL'] + df_trades.loc[i,'GAIN']

优化方案（非循环实现） #

利用Pandas的分组、移位(shift)、累计求和(cumsum)等向量化操作替代循环，步骤如下：

1. 标记交易组：对TRADE==1的行做累计求和，将每一组买卖（买+卖）归为同一组
2. 计算分组内的关键值：提取每组的买入价、买入时的资金，计算持仓量
3. 填充各字段值：通过向量化操作批量计算DIFF、GAIN、GAIN_C、CAPITAL和POSITION

优化后代码：

import pandas as pd

# 1. 标记交易组：每一次买入开始一个新组
df_trades['GROUP'] = (df_trades['TRADE'] == 1).cumsum()

# 2. 提取每组的买入价、买入时的初始资金
group_info = df_trades[df_trades['TRADE'] == 1].set_index('GROUP')[['<Open>', 'CAPITAL']].rename(columns={'<Open>': 'BUY_PRICE', 'CAPITAL': 'BUY_CAPITAL'})
df_trades = df_trades.merge(group_info, on='GROUP', how='left')

# 3. 计算POSITION：买入行用资金整除买入价，卖出行设为0
df_trades['POSITION'] = df_trades.apply(lambda x: x['BUY_CAPITAL'] // x['BUY_PRICE'] if x['TRADE'] == 1 else 0, axis=1)

# 4. 计算DIFF：仅卖出行计算当前价与买入价的差值，其余为NaN
df_trades['DIFF'] = df_trades.apply(lambda x: x['<Open>'] - x['BUY_PRICE'] if x['TRADE'] == -1 else pd.NA, axis=1)

# 5. 计算GAIN：卖出行用差值乘对应持仓量，其余为0
df_trades['GAIN'] = df_trades.apply(lambda x: x['DIFF'] * x['POSITION'] if x['TRADE'] == -1 else 0, axis=1)

# 6. 计算累计收益GAIN_C：对GAIN做累计求和
df_trades['GAIN_C'] = df_trades['GAIN'].cumsum()

# 7. 计算CAPITAL：初始资金加累计收益，买入行继承上一行卖出后的资金
df_trades['CAPITAL'] = df_trades['CAPITAL'].iloc[0] + df_trades['GAIN_C']
buy_mask = df_trades['TRADE'] == 1
df_trades.loc[buy_mask, 'CAPITAL'] = df_trades['CAPITAL'].shift(1).loc[buy_mask]

# 清理临时字段并调整列顺序
df_trades.drop(['GROUP', 'BUY_PRICE', 'BUY_CAPITAL'], axis=1, inplace=True)
df_trades = df_trades[['<Open>', 'TRADE', 'CAPITAL', 'POSITION', 'GAIN', 'GAIN_C', 'DIFF']]

补充说明 #

• 该方案假设交易为严格买-卖交替（TRADE序列为1,-1,1,-1...），若存在连续买入/卖出场景，可调整分组逻辑适配
• 向量化操作避免了逐行循环，数据量越大性能提升越明显
• 最终输出格式与示例完全匹配

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