如何将Top3预测结果与概率得分作为不同列添加到pandas DataFrame

实现代码 #

import pandas as pd
import numpy as np
# 创建DataFrame
data = [
    ["id_1",6,7,9],
    ["id_2",9,7,1],
    ["id_3",3,5,5],
]
test = pd.DataFrame(data, columns = ['Student_Id', 'Math', 'Physical','Arts'])
# 来自clf.Classes_的分类类别
Classes = np.array(['id_1', 'id_2', 'id_3'])
# 来自predict_proba()的概率结果
top_test_probabilities = np.array([[0.70, 0.30, 0.10], [0.30, 0.70, 0.10], [0.10, 0.30, 0.70]])
# 排序后的Top3值对应的索引（按概率从高到低排序）
best_3n = np.array([[0, 1, 2], [1, 0, 2], [2, 1, 0]])

# 直接给原DataFrame赋值新增预测ID列
test[["Predicted_id_1", "Predicted_id_2", "Predicted_id_3"]] = Classes[best_3n]
# 按行取排序后的对应概率，赋值新增概率列
test[["Probability_Score_1", "Probability_Score_2", "Probability_Score_3"]] = np.take_along_axis(top_test_probabilities, best_3n, axis=1)

执行完成后test变量就是符合要求的输出结果。

逻辑说明 #

• 直接使用pandas多列赋值特性，不需要额外生成中间DataFrame再拼接，代码更简洁
• np.take_along_axis按行（axis=1）匹配best_3n的索引取出对应概率，和排序后的ID顺序完全对应
• 原数据和新增列的行数完全对齐，不会出现顺序错乱问题

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