推薦系統測評指標——計算DCG、IDCG以及nDCG的python代碼

1. 公式

DCG： D C G @ K = ∑ i = 1 K 2 r i ? 1 l o g 2 ( i + 1 ) DCG@K=\sum_{i=1}^{K}\frac{2^{r_i}-1}{log_2(i+1)} DCG@K=i=1∑K?log2?(i+1)2ri??1?其中， K K K是推薦串列的大小；
i i i 是指推薦串列中的第 i i i個推薦項；
r i r_i ri?是指推薦項 i i i 是否被用戶點擊，若點擊則為 1 1 1，否則為 0 0 0，在實際測驗中，我們通常吧推薦串列中在測驗集的推薦項視為被用戶點擊的推薦項，不在測驗集中的推薦項視為未被用戶點擊的推薦項；

IDCG： I D C G @ K = ∑ i = 1 K 2 r i ? 1 l o g 2 ( i + 1 ) IDCG@K=\sum_{i=1}^{K}\frac{2^{r_i}-1}{log_2(i+1)} IDCG@K=i=1∑K?log2?(i+1)2ri??1?注意到IDCG的計算公式與DCG一致，區別在于IDCG是完美的DCG，也即 r i = 1 r_i=1 ri?=1的推薦項在推薦串列 K K K頭部， r i = 0 r_i=0 ri?=0的推薦項在推薦串列 K K K的末尾，

nDCG： n D C G @ K = D C G I D C G nDCG@K=\frac{DCG}{IDCG} nDCG@K=IDCGDCG?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)