关于sklearn中ndcg_score输入格式及多查询场景使用的问询

我刚接触这个函数的时候也懵过，结合你已经算出多个查询排序概率的场景，给你掰扯清楚核心要点：

核心维度规则 #

ndcg_score是按查询（query）维度独立计算的，所以输入的两个参数必须满足：

• y_true：二维数组，形状为(n_queries, n_samples_per_query)。每个子数组对应一个查询下，所有候选样本的真实相关性得分（比如用0、1、2表示不相关、相关、非常相关，数字越大相关性越强）。
• y_score：和y_true形状完全匹配的二维数组，每个子数组对应同一个查询下，所有候选样本的预测排序概率/得分（得分越高，你的模型认为这个样本越相关）。

对应官方示例的解释 #

官方给的示例：

>>> y_true = [1, 0, 2]
>>> y_score = [[0.15, 0.55, 0.2]]
>>> ndcg_score(y_true, y_score)
0.650000...

这里的y_true = [1,0,2]其实是单个查询的情况，sklearn会自动把它转换成形状为(1,3)的二维数组；y_score的二维数组也对应这单个查询的3个候选样本的预测得分。

你的多查询场景示例 #

假设你有2个查询，每个查询对应不同数量的候选样本，直接看代码更直观：

from sklearn.metrics import ndcg_score

# 2个查询的真实相关性得分
# 查询1：3个候选样本，真实得分[2,0,1] → 第1个样本最相关，第3个次之，第2个最不相关
# 查询2：4个候选样本，真实得分[3,1,0,2] → 第1个最相关，第4个次之，第2个，第3个最不相关
y_true = [
    [2, 0, 1],
    [3, 1, 0, 2]
]

# 对应每个查询的预测排序概率（得分越高，模型认为越相关）
# 查询1的预测：模型觉得第2个样本最相关（0.6），第1个次之（0.3），第3个最不相关（0.1）
# 查询2的预测：模型觉得第2个样本最相关（0.7），第4个次之（0.5），第1个（0.2），第3个最不相关（0.1）
y_score = [
    [0.3, 0.6, 0.1],
    [0.2, 0.7, 0.1, 0.5]
]

# 计算nDCG，默认返回所有查询的nDCG平均值
average_ndcg = ndcg_score(y_true, y_score)
print(average_ndcg)

避坑小提示 #

• 千万别把所有查询的得分拼成一维数组！必须按查询分组，每个查询对应一个子数组。
• 真实得分y_true的数值是相关性的量化等级，只要保证"数字越大越相关"就行，不用局限于0/1，可以是0-5这类多等级得分。
• 每个查询的候选样本数量可以不一样，只要每个查询对应的y_true子数组和y_score子数组长度一致就ok。
• 不需要提前对预测得分排序，ndcg_score会自动根据y_score的得分高低对样本排序，再计算这个排序和真实最优排序的差距。

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