如何从XGBoost的predict_proba()生成数组提取类别1概率≥90%的元素？

首先明确opt.predict_proba(x_test)的输出结构：返回的是二维数组，每一行对应一个测试样本，列对应不同类别（假设类别顺序为[0,1]），所以类别1的概率是每行的第2个元素（索引为1）。你之前直接和[0.9, 0.1]比较的方式错误，因为数组逐元素比较需要维度匹配，且这并非筛选类别1概率的正确逻辑。以下是正确实现方法：

方法一：使用numpy索引（推荐） #

假设已导入numpy：

import numpy as np

# 获取所有测试样本的预测概率数组
probs = opt.predict_proba(x_test)

# 筛选出类别1概率≥0.9的样本概率数据
high_confidence_probs = probs[probs[:, 1] >= 0.9]

# 若需要同时获取对应的测试样本数据：
high_confidence_samples = x_test[probs[:, 1] >= 0.9]

方法二：使用列表推导式 #

如果不依赖numpy，也可以通过遍历实现：

probs = opt.predict_proba(x_test)

# 筛选符合条件的概率行
high_confidence_probs = [prob for prob in probs if prob[1] >= 0.9]

关键逻辑说明 #

• probs[:, 1]会提取二维数组中所有行的第2列数据（即所有样本的类别1概率），生成一维布尔数组
• 用probs[:, 1] >= 0.9得到的布尔数组作为索引，即可筛选出原概率数组中符合条件的行

内容的提问来源于stack exchange，提问作者Eduardo Schneider Arruda