如何高效替代Pandas DataFrame中遍历列值的三重循环？

问题背景 #

给定如下结构的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({'Level': ['a','b', 'c'], 'Kontogruppe': ['a', 'a', 'b'], 
                   'model': ["alpha", "beta", "alpha"], 'MSE': [0, 1 ,1],
                   'actual_value': [1,2,3], 'forecast_value': [2,2,2]}) 

原处理逻辑通过三重循环遍历Level、Kontogruppe、model的唯一值，筛选子数据集计算out_MSE（实际值总和/预测值总和，空数据集则设为10^10），最后合并回原DataFrame，但数据量大时循环性能极差：

def metrics(df):
    df_map= pd.DataFrame({'Level': ['a'], 'Kontogruppe': ['a'],
                            'model': ["alpha"], 'MSE': [0]})
    for i in df['Level'].unique():
        for j in df['Kontogruppe'].unique():
            for k in df['model'].unique():
                df_lkm = df.loc[(df['Level'] == i) & (df['Kontogruppe'] == j) & 
                                    (df['model'] == k)]
                if df_lkm.empty:
                    out_MSE = 10000000000                   

                else:
                    out_MSE = sum(df_lkm['actual_value'])/sum(df_lkm['forecast_value'])                    

                df_map_map = pd.DataFrame({'Level': [i], 'Kontogruppe': [j], 'model': [k], 
                                        'out_MSE': [out_MSE]}) 
                df_map = pd.concat([df_map, df_map_map])
            
    df = pd.merge(df, df_map, how='left', on=['Level', 'Kontogruppe', 'model'])   
                    
    return df

df = metrics(df)

高效解决方案 #

可以利用Pandas的分组聚合+笛卡尔积生成实现完全矢量化操作，性能比循环提升显著：

步骤1：分组计算有效组合的out_MSE #

通过groupby按三个维度分组，直接计算每组的实际值总和与预测值总和的比值：

# 分组计算实际值总和/预测值总和
grouped = df.groupby(['Level', 'Kontogruppe', 'model']).apply(
    lambda x: x['actual_value'].sum() / x['forecast_value'].sum()
).reset_index(name='out_MSE')

步骤2：生成所有维度的笛卡尔积 #

用itertools.product生成三个维度唯一值的所有组合，确保覆盖原数据中不存在的空组合：

from itertools import product

# 获取各维度的唯一值集合
levels = df['Level'].unique()
kontos = df['Kontogruppe'].unique()
models = df['model'].unique()

# 生成所有可能的组合
all_combinations = pd.DataFrame(product(levels, kontos, models), 
                                columns=['Level', 'Kontogruppe', 'model'])

步骤3：合并并填充空组合的默认值 #

将分组结果与笛卡尔积合并，空组合的out_MSE填充为10^10，最后合并回原DataFrame：

# 合并分组结果与所有组合，填充缺失值
df_map = all_combinations.merge(grouped, on=['Level', 'Kontogruppe', 'model'], how='left')
df_map['out_MSE'] = df_map['out_MSE'].fillna(10000000000)

# 合并回原DataFrame
df_final = df.merge(df_map, on=['Level', 'Kontogruppe', 'model'], how='left')

完整优化函数 #

def metrics_optimized(df):
    # 分组计算有效组合的out_MSE
    grouped = df.groupby(['Level', 'Kontogruppe', 'model']).apply(
        lambda x: x['actual_value'].sum() / x['forecast_value'].sum()
    ).reset_index(name='out_MSE')
    
    # 生成所有维度的笛卡尔积
    from itertools import product
    all_combinations = pd.DataFrame(
        product(df['Level'].unique(), df['Kontogruppe'].unique(), df['model'].unique()),
        columns=['Level', 'Kontogruppe', 'model']
    )
    
    # 合并并填充空组合的默认值
    df_map = all_combinations.merge(grouped, on=['Level', 'Kontogruppe', 'model'], how='left')
    df_map['out_MSE'] = df_map['out_MSE'].fillna(10000000000)
    
    # 合并回原DataFrame
    return df.merge(df_map, on=['Level', 'Kontogruppe', 'model'], how='left')

# 使用优化后的函数处理数据
df = metrics_optimized(df)

优化说明 #

• groupby是Pandas基于C实现的矢量化操作，比Python层级的循环快几个数量级；
• itertools.product生成笛卡尔积的效率远高于嵌套循环遍历；
• 用merge+fillna处理空组合逻辑，代码更简洁且性能稳定。

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