Pandas比较多个等长DataFrame 保留最大概率及对应分类

现有4个结构、行数完全一致的DataFrame：df1、df2、df3、df4，均包含max_proba（预测最大概率）、chosen_class（对应预测分类）两列，需逐行对比4个表的概率值，提取每行最大概率及对应的分类结果。
示例表结构如下：

df1
   max_proba    chosen_class
0   0.8            class_A
1   0.92           class_B
2   0.82           class_B
3   0.74           class_B
4   0.58           class_A

df2
   max_proba    chosen_class
0   0.6            class_C
1   0.62           class_D
2   0.87           class_D
3   0.94           class_C
4   0.62           class_D

# df3、df4结构同上，仅概率、分类取值存在差异

实现代码 #

方法1：pandas原生实现（可读性强，易排查问题） #

import pandas as pd

# 第一步：统一重置索引，避免原索引不连续/重复导致行错位
df_list = [df.reset_index(drop=True) for df in [df1, df2, df3, df4]]

# 第二步：拼接所有表的max_proba列，生成每行4个概率值的矩阵
proba_cols = pd.concat([df["max_proba"] for df in df_list], axis=1)

# 第三步：逐行定位最大值所在的列位置（0对应df1，1对应df2，以此类推）
row_max_col = proba_cols.idxmax(axis=1)

# 第四步：生成结果表
result = pd.DataFrame()
result["max_proba"] = [proba_cols.loc[row_idx, col_idx] for row_idx, col_idx in enumerate(row_max_col)]
result["chosen_class"] = [df_list[col_idx].loc[row_idx, "chosen_class"] for row_idx, col_idx in enumerate(row_max_col)]

方法2：numpy实现（数据量大时运算速度更快） #

import numpy as np
import pandas as pd

df_list = [df.reset_index(drop=True) for df in [df1, df2, df3, df4]]

# 转成numpy数组提升运算速度
proba_array = np.array([df["max_proba"].values for df in df_list]).T
class_array = np.array([df["chosen_class"].values for df in df_list]).T

# 逐行找最大值位置
max_pos = proba_array.argmax(axis=1)
row_index = np.arange(len(max_pos))

result = pd.DataFrame({
    "max_proba": proba_array[row_index, max_pos],
    "chosen_class": class_array[row_index, max_pos]
})

逻辑验证 #

以给出的示例行为准：某行df1概率0.23、df2概率0.86、df3概率0.56、df4概率0.76，代码会识别到最大值在第1列（对应df2），最终返回概率值0.86、对应分类class_E，完全符合需求。

注意事项 #

• 若4个表的样本顺序无法保证完全一致，不要直接按位置拼接，需先通过唯一样本ID做表关联对齐后再计算，避免结果错配
• 若同一行存在多个表的max_proba值完全相等的情况，上述逻辑默认返回排在列表最前面的表对应的分类结果

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