Python Pandas回撤计算：向量化改写后结果不符的问题排查

问题背景 #

原函数需要对长度递增的DataFrame子序列取最小值、长度递减的子序列取最大值，计算对应的回撤值后取中位数。原实现采用三层嵌套循环，时间复杂度极高；尝试用cummax/cummin向量化改写后，结果与原循环版本不一致，需要排查错误并给出正确优化方案。

原循环代码（正确逻辑） #

def drawdown(result_df, dict_dfs):
    last_year_df = pd.DataFrame(data=np.nan, index=result_df.index, columns=result_df.columns)

    for idx in range(len(result_df)):
        for stock in result_df.columns:
            past_date_idx = max(0, idx - 250)
            past_date = result_df.index[past_date_idx]
            current_date = result_df.index[idx]

            last_year = dict_dfs['close'].loc[past_date:current_date, stock]

            drawdowns = []
            for i in range(len(last_year)):
                rolling_min = last_year.iloc[:i + 1].min()
                rolling_max = last_year.iloc[i:].max()
                if rolling_min != 0:
                    drawdown = (rolling_max - rolling_min) / rolling_min
                    drawdowns.append(drawdown)

            last_year_df.iloc[idx][stock] = np.median(drawdowns)

    return last_year_df

错误的向量化改写代码 #

def drawdown(result_df, dict_dfs):
    last_year_df = pd.DataFrame(data=np.nan, index=result_df.index, columns=result_df.columns)

    for idx in range(len(result_df)):
        for stock in result_df.columns:
            past_date_idx = max(0, idx - 250)
            past_date = result_df.index[past_date_idx]
            current_date = result_df.index[idx]

            last_year = dict_dfs['close'].loc[past_date:current_date, stock]
            max_dataframe = last_year.iloc[::-1]
            min_dataframe = last_year

            rolling_max = max_dataframe.cummax()
            rolling_min = min_dataframe.cummin()

            drawdown = (rolling_max - rolling_min) / rolling_min

            last_year_df.loc[current_date] = drawdown.iloc[-1]

    return last_year_df

错误原因分析 #

1. 最大值计算逻辑错误
   原代码中rolling_max是从第i个位置到序列末尾的全局最大值，而改写代码仅对反转后的序列取cummax，未将结果反转回原顺序，导致rolling_max的位置对应关系完全错误。

2. 未实现中位数计算
   原代码是对每个位置i计算的回撤值取中位数，而改写代码直接取了回撤序列的最后一个值，完全偏离原逻辑。

3. 列赋值错误
   last_year_df.loc[current_date] = drawdown.iloc[-1]会将当前日期的所有股票列覆盖为同一个值，与原代码逐个股票赋值的逻辑不符。

正确的向量化优化方案 #

基础修正版（修复逻辑错误，保留外层循环） #

import pandas as pd
import numpy as np

def drawdown(result_df, dict_dfs):
    last_year_df = pd.DataFrame(data=np.nan, index=result_df.index, columns=result_df.columns)
    
    for idx in range(len(result_df)):
        for stock in result_df.columns:
            past_date_idx = max(0, idx - 250)
            past_date = result_df.index[past_date_idx]
            current_date = result_df.index[idx]
            
            last_year = dict_dfs['close'].loc[past_date:current_date, stock].dropna()
            if len(last_year) == 0:
                continue
            
            # 计算从开头到每个位置的累计最小值（对应原代码的rolling_min）
            cum_min = last_year.cummin()
            # 计算从每个位置到末尾的最大值：反转后取cummax再反转回原顺序
            reverse_cum_max = last_year.iloc[::-1].cummax().iloc[::-1]
            
            # 过滤无效值（避免除以0）并计算回撤
            valid_mask = cum_min != 0
            drawdowns = (reverse_cum_max[valid_mask] - cum_min[valid_mask]) / cum_min[valid_mask]
            
            if not drawdowns.empty:
                last_year_df.loc[current_date, stock] = np.median(drawdowns)
    
    return last_year_df

进阶优化版（用滚动窗口减少循环） #

利用rolling窗口批量处理每个股票的所有日期，进一步降低时间复杂度：

def drawdown_rolling_optimized(result_df, dict_dfs):
    close_df = dict_dfs['close'].reindex(result_df.index)
    last_year_df = pd.DataFrame(index=result_df.index, columns=result_df.columns)
    
    def calculate_window_median(window):
        # 窗口内计算累计最小值和从每个位置到末尾的最大值
        cum_min = window.cummin()
        reverse_cum_max = window[::-1].cummax()[::-1]
        # 过滤无效值并计算回撤
        valid_mask = cum_min != 0
        drawdowns = (reverse_cum_max[valid_mask] - cum_min[valid_mask]) / cum_min[valid_mask]
        return np.median(drawdowns) if not drawdowns.empty else np.nan
    
    # 对每个股票应用251个交易日的滚动窗口（包含当前日期，往前250天）
    for stock in result_df.columns:
        last_year_df[stock] = close_df[stock].rolling(window=251, min_periods=1).apply(calculate_window_median)
    
    return last_year_df

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