推荐系统测评指标——计算DCGIDCG以及nDCG的python代码

Posted 2021-09-06 白水baishui

1. 公式

DCG：$$DCG@K=\sum _{i=1}^K\frac{2^{r_i}-1}{lo{g}_2(i+1)}$$其中，$K$是推荐列表的大小；
$i$ 是指推荐列表中的第$i$个推荐项；
$r_i$是指推荐项 $i$ 是否被用户点击，若点击则为$1$，否则为$0$，在实际测试中，我们通常吧推荐列表中在测试集的推荐项视为被用户点击的推荐项，不在测试集中的推荐项视为未被用户点击的推荐项；

IDCG：$$IDCG@K=\sum _{i=1}^K\frac{2^{r_i}-1}{lo{g}_2(i+1)}$$注意到IDCG的计算公式与DCG一致，区别在于IDCG是完美的DCG，也即$r_i=1$的推荐项在推荐列表$K$头部，$r_i=0$的推荐项在推荐列表$K$的末尾。

nDCG：$$nDCG@K=\frac{DCG}{IDCG}$$NDCG是归一化后的DCG

2. 代码

import numpy as np

def DCG(A, test_set):
    # ------ 计算 DCG ------ #
    dcg = 0
    for i in range(len(A)):
        # 给r_i赋值，若r_i在测试集中则为1，否则为0
        r_i = 0
        if A[i] in test_set:
            r_i = 1
        dcg += (2 ** r_i - 1) / np.log2((i + 1) + 1) # (i+1)是因为下标从0开始
    return dcg

def IDCG(A, test_set):
    # ------ 将在测试中的a排到前面去，然后再计算DCG ------ #
    A_temp_1 = [] # 临时A，用于存储r_i为1的a
    A_temp_0 = []  # 临时A，用于存储r_i为0的a
    for a in A:
        if a in test_set:
            # 若a在测试集中则追加到A_temp_1中
            A_temp_1.append(a)
        else:
            # 若a不在测试集中则追加到A_temp_0中
            A_temp_0.append(a)
    A_temp_1.extend(A_temp_0)
    idcg = DCG(A_temp_1, test_set)
    return idcg

def NDCG(A, test_set):
    dcg = DCG(A, test_set) # 计算DCG
    idcg = IDCG(A, test_set) # 计算IDCG
    if dcg == 0 or idcg == 0:
        ndcg = 0
    else:
        ndcg = dcg / idcg
    return ndcg

if __name__ == "__main__":
	# ------ 计算推荐列表A的NDCG ------ #
	# A：推荐列表，一维list，存储了推荐算法推荐出的推荐项的id
	# test_set：测试集，一维list，存储了测试集推荐项的id
    ndcg_A = NDCG(A, test_set)