如何优化pandas DataFrame嵌套循环计算交易PnL的运行速度？

优化思路 #

你当前的嵌套循环是O(n²)时间复杂度，146万行数据遇到7000个买点时最坏要执行近10亿次空值判断，所以运行极慢。我们可以通过双指针法把时间复杂度降到O(n)，全程只需遍历买卖点各一次，几秒就能跑完。

具体实现步骤 #

1. 提前提取所有非空买点、卖点的索引，避免逐行判断空值
2. 预先生成下一根K线的Open列，避免计算时重复做索引偏移
3. 用单指针遍历卖点数组，给每个买点匹配第一个出现的后续卖点

优化代码

import pandas as pd
import numpy as np

# 预计算下一根K线的Open价
dfcondensed['next_open'] = dfcondensed['Open'].shift(-1)
# 提取所有非空买卖点的索引
buy_idx_list = dfcondensed[~dfcondensed['Buy'].isna()].index
sell_idx_list = dfcondensed[~dfcondensed['Sell'].isna()].index

pnl_result = []
sell_ptr = 0
max_sell_pos = len(sell_idx_list)

# 双指针匹配买卖点
for buy_pos in buy_idx_list:
    # 找到第一个在当前买点之后的卖点
    while sell_ptr < max_sell_pos and sell_idx_list[sell_ptr] <= buy_pos:
        sell_ptr += 1
    # 后续无可用卖点直接终止
    if sell_ptr >= max_sell_pos:
        break
    # 计算PnL
    buy_entry_price = dfcondensed.at[buy_pos, 'next_open']
    sell_exit_price = dfcondensed.at[sell_idx_list[sell_ptr], 'next_open']
    if pd.notna(buy_entry_price) and pd.notna(sell_exit_price):
        pnl = (sell_exit_price - buy_entry_price) / buy_entry_price * 100
        pnl_result.append(pnl)

极致加速版本（可选）

如果需要更快的运行速度，可以把数据转成numpy数组操作，索引效率更高：

# 转numpy数组加速
next_open_arr = dfcondensed['next_open'].values
buy_idx_arr = buy_idx_list.values
sell_idx_arr = sell_idx_list.values

sell_ptr = 0
max_sell_pos = len(sell_idx_arr)
pnl_result = []

for buy_pos in buy_idx_arr:
    while sell_ptr < max_sell_pos and sell_idx_arr[sell_ptr] <= buy_pos:
        sell_ptr += 1
    if sell_ptr >= max_sell_pos:
        break
    buy_price = next_open_arr[buy_pos]
    sell_price = next_open_arr[sell_idx_arr[sell_ptr]]
    if not np.isnan(buy_price) and not np.isnan(sell_price):
        pnl_result.append((sell_price - buy_price)/buy_price * 100)

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