DataFrame分组计算胜率与期望的代码问题及优化咨询

你遇到的核心问题是**df2的结构和stats函数的调用逻辑不匹配**：

1. df2是通过groupby(['day', 'symbol', 'strategy', 'isWinner'])聚合得到的，此时isWinner是df2的索引层级，而非列名，所以在后续分组后的stats函数里，无法通过s['isWinner']访问这个字段。
2. 你的stats函数没有返回值，就算逻辑正确也无法生成有效结果。

首先调整stats函数，利用分组后的索引（isWinner的True/False值）获取胜负组的统计数据，同时返回胜率和期望：

import pandas as pd

def stats(s):
    # 从索引中提取胜负组的统计值，处理无胜负组的边界情况
    win_count = s.loc[True, ('pnl', 'count')] if True in s.index else 0
    lose_count = s.loc[False, ('pnl', 'count')] if False in s.index else 0
    total_count = win_count + lose_count
    
    winrate = win_count / total_count if total_count > 0 else 0
    
    win_mean = s.loc[True, ('pnl', 'mean')] if True in s.index else 0
    lose_mean = s.loc[False, ('pnl', 'mean')] if False in s.index else 0
    expectancy = win_mean * winrate - lose_mean * (1.0 - winrate)
    
    return pd.Series({'winrate': winrate, 'expectancy': expectancy})

# 示例数据初始化（可替换为你的真实数据）
df = pd.DataFrame({
    'day': ['2024-01-01']*10 + ['2024-01-02']*10,
    'symbol': ['AAPL']*5 + ['MSFT']*5 + ['AAPL']*5 + ['MSFT']*5,
    'strategy': ['strat1']*3 + ['strat2']*2 + ['strat1']*3 + ['strat2']*2 + ['strat1']*3 + ['strat2']*2,
    'pnl': [10, -5, 15, -3, 8, 20, -10, 5, -2, 12, -8, 14, -6, 9, -1, 18, -7, 11, -4, 16]
})

df["isWinner"] = df['pnl'] >= 0
df2 = df.groupby(['day', 'symbol', 'strategy', 'isWinner']).agg({'pnl': ['count', 'mean', 'std', 'min', 'max']})

# 计算胜率和期望
stats_df = df2.groupby(['day', 'symbol', 'strategy']).apply(stats)

由于stats_df是按day、symbol、strategy分组的结果，我们可以用join方法将其合并到df2中：

# 重置df2索引以便合并，之后恢复原索引结构
df2 = df2.reset_index().join(stats_df, on=['day', 'symbol', 'strategy']).set_index(['day', 'symbol', 'strategy', 'isWinner'])

当前实现需要两次分组操作，我们可以简化流程，避免创建辅助列，直接在一次分组中计算所有指标：

def calculate_stats(group):
    # 计算胜负次数与均值
    win_mask = group['pnl'] >= 0
    win_count = win_mask.sum()
    lose_count = len(group) - win_count
    total_count = win_count + lose_count
    
    winrate = win_count / total_count if total_count > 0 else 0
    
    win_mean = group.loc[win_mask, 'pnl'].mean() if win_count > 0 else 0
    lose_mean = group.loc[~win_mask, 'pnl'].mean() if lose_count > 0 else 0
    expectancy = win_mean * winrate - lose_mean * (1.0 - winrate)
    
    # 计算pnl的其他统计量
    pnl_stats = group['pnl'].agg(['count', 'mean', 'std', 'min', 'max'])
    
    # 合并所有结果
    result = pnl_stats.to_dict()
    result.update({'winrate': winrate, 'expectancy': expectancy})
    return pd.Series(result)

# 直接按目标维度分组计算所有指标
final_df = df.groupby(['day', 'symbol', 'strategy']).apply(calculate_stats).reset_index()

该方案的优势：

• 无需额外创建isWinner等辅助列，减少内存占用
• 仅需一次分组操作，计算效率更高
• 结果直接整合所有统计量，无需后续合并步骤

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