Python及PySpark中提取分类器最大得分的最优方法

你可以通过pandas的分组筛选功能，按Observation分组后提取每组中Probability最大的行，就能得到每个观测对应的获胜类别及概率。

代码示例 #

假设你的DataFrame名为df，以下是具体实现：

import pandas as pd

# 构造原始数据（也可以是你读取的已有数据集）
data = {
    'Observation': [1,1,2,2,3,3],
    'Class': [0,1,0,1,0,1],
    'Probability': [0.5013, 0.4987, 0.5010, 0.4990, 0.5128, 0.4872]
}
df = pd.DataFrame(data)

# 方法1：分组后筛选最大概率行
result_df = df.groupby('Observation', as_index=False).apply(
    lambda x: x[x['Probability'] == x['Probability'].max()]
).reset_index(drop=True)

# 方法2：更高效的写法（适合大数据量）
# result_df = df.loc[df.groupby('Observation')['Probability'].idxmax()].reset_index(drop=True)

print(result_df)

输出结果 #

说明 #

• 方法1通过分组后对每组进行条件筛选，逻辑直观易懂
• 方法2利用idxmax()直接获取每组最大概率值的索引，再提取对应行，执行效率更高，适合处理大规模数据

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