NDCG 计算

NDCG 计算题解

题意

给定 JSON 数据，包含：

• y_true：每个样本的真实相关性标签；
• y_score：模型给出的预测分数；
• K：只计算前 K 个结果的指标。

需要计算并输出：

• dcg：按预测分数降序排序后的折损累计增益；
• idcg：按真实相关性降序排序后的理想折损累计增益；
• ndcg：归一化折损累计增益，即

$$\mathrm{NDCG} = \frac{\mathrm{DCG}}{\mathrm{IDCG}}$$

最终以紧凑 JSON 格式输出。

计算公式

对于排序后的前 $K$ 个结果，第 $r$ 个位置的真实相关性记为 $rel_r$，则：

$$\mathrm{DCG@K} = \sum_{r=1}^{K} \frac{2^{rel_r}-1}{\log_2(r+1)}$$

理想排序是将所有样本按真实相关性从大到小排序，此时得到：

$$\mathrm{IDCG@K} = \sum_{r=1}^{K} \frac{2^{rel_r^{ideal}}-1}{\log_2(r+1)}$$

因此：

$$\mathrm{NDCG@K} = \frac{\mathrm{DCG@K}}{\mathrm{IDCG@K}}$$

若 $\mathrm{IDCG@K}=0$，则令 $\mathrm{NDCG@K}=0$。

思路

1. 读取 JSON 字符串，解析出 y_true、y_score、K。

2. 将每个样本组织为三元组：

   $$(i,\ y\_true_i,\ y\_score_i)$$

3. 按预测分数从大到小排序。若分数相同，则按下标从小到大排序，保证结果稳定。

4. 取排序后的前 $K$ 个样本，按照 DCG 公式累加，得到 dcg。

5. 再按真实相关性从大到小排序。若相关性相同，则按下标从小到大排序，得到理想排序。

6. 取前 $K$ 个样本，按照同样公式累加，得到 idcg。

7. 计算 ndcg = dcg / idcg，若 idcg == 0，则 ndcg = 0。

8. 将 idcg 和 ndcg 保留 6 位小数，最后用紧凑 JSON 输出。

正确性说明

DCG 的定义是对排序结果中每个位置的相关性进行指数加权，并根据位置进行对数折损。按预测分数排序得到实际排序，因此计算出的就是模型排序下的 DCG。

IDCG 是同一批样本在最优排序下的 DCG。由于指数函数 $2^x$ 单调递增，将真实相关性从大到小排列即可使 DCG 最大，因此该排序就是理想排序。

NDCG 用 IDCG 对 DCG 归一化，使得不同查询或不同样本集合之间的指标可以比较。

复杂度分析

设样本数量为 $n$。

• 排序两次，时间复杂度为：

$$O(n\log n)$$

• 计算 DCG 和 IDCG 各遍历前 $K$ 个元素，时间复杂度为：

$$O(K)$$

因此总时间复杂度为：

$$O(n\log n)$$

空间复杂度为：

$$O(n)$$

参考代码

import json
import numpy

s = input().strip()
data = json.loads(s)

ytrue = data['y_true']
yscore = data['y_score']
k = data['K']
n = len(ytrue)

arr = []
for i in range(n):
    arr.append((i, ytrue[i], yscore[i]))

# 按预测分数降序排序，同分按下标升序
arr.sort(key=lambda x: (-x[2], x[0]))

dcg = 0
for r in range(1, k + 1):
    dcg += (2.0 ** arr[r - 1][1] - 1) / numpy.log2(r + 1.0)

# 按真实相关性降序排序，同分按下标升序
arr.sort(key=lambda x: (-x[1], x[0]))

idcg = 0
for r in range(1, k + 1):
    idcg += (2.0 ** arr[r - 1][1] - 1) / numpy.log2(r + 1.0)

ndcg = 0
if idcg != 0:
    ndcg = dcg / idcg

dcg = round(dcg, 6)
idcg = round(idcg, 6)
ndcg = round(ndcg, 6)

data = {
    "dcg": dcg,
    "idcg": idcg,
    "ndcg": ndcg
}

print(json.dumps(data, separators=(',', ':')))

输入示例

{"y_true":[3,2,0,1],"y_score":[0.2,0.8,0.8,0.1],"K":3}

输出示例

{"dcg":6.5,"idcg":9.392789,"ndcg":0.69202}

注意事项

1. 排序时需要保证同分情况下结果稳定，因此代码中使用下标作为第二排序关键字。

2. 题目通常保证 $K \le n$，否则取前 $K$ 个元素时可能越界。

3. idcg 为 $0$ 时不能直接做除法，需要单独处理。

4. 输出时使用：

   json.dumps(data, separators=(',', ':'))

   可以避免 JSON 中出现多余空格。