信息检索————NDCG计算

先贴代码，原理有时间补上。

import numpy as np
import copy
from math import log

# input data
ideal = np.array([5,5,4,4,3,3,2,2,1,1])
relevant_score_list = np.array([5.0,5.0,4.0,4.0,3.0,3.0,2.0,2.0,1.0,1.0])
recall_list = np.array([8,11,2,12,3,1,5,4,13,7])

def dcg(rscore, m):
    return (2.0 ** rscore - 1.0) / np.log2(2.0 + m)

def dcg_k(rslist,rlist,k):
    # 计算第k个dcg值
    dcgscore_list = [0.0 for i in range(11)]##计算1,2,...,k的dcg值
    for i in range(k):
        relevant_score = 0
        if rlist[i] < len(rslist):
            relevant_score = rslist[rlist[i]-1]
        dcgscore_list[i+1] = dcgscore_list[i] + dcg(relevant_score, i)
    return dcgscore_list;

def idcg_k(ideal,k):
    # idcg是一个系统最理想情况下的返回结果排序，也就是一个案例中所给的最佳排序
    idcgscore_list = [0.0 for i in range(11)] #计算1,2,...,k的idcg值
    for i in range(k):
        idcgscore_list[i+1] = idcgscore_list[i] + dcg(ideal[i], i)
    return idcgscore_list;

def ndcg(rslist,rlist,k):
    dcgscore_list = dcg_k(rslist,rlist,k)
    # 计算归一化因子z, 最完美情况应该是recall结果按照相关性分数降序排列
    idcgscore_list = idcg_k(ideal,k)

    ndcg_list = [0.0 for i in range(k)]
    for i in range(k):
        ndcg_list[i] = round(dcgscore_list[i+1] / idcgscore_list[i+1], 3)
    print(ndcg_list)
    

def ndcg(ideal, rlist, k):
    recall = [0.0 for i in range(k)]
    for i in range(k):
        if rlist[i] <= len(ideal):
            recall[i] = ideal[rlist[i]-1]
    
    dcg, dcg_max, ndcg = 0.0, 0.0, 0.0
    for i, (true, predict) in enumerate(zip(ideal, recall)):
        dcg += (2 ** predict - 1) / log(2 + i)
        dcg_max += (2 ** true -1) / log(2 + i)
        ndcg = dcg / dcg_max
    return ndcg
    
# print(ndcg(relevant_score_list,recall_list,10))
# ndcg(relevant_score_list,recall_list,10)
ndcg(ideal, recall_list, 10)

代码中两种方式实现了ndcg的计算，调用的时候需要注意一下。