CG, DCG, IDCG, NDCG的实现

import numpy as np


class Solution(object):
    def __init__(self, relevance, p=6):
        """
        https://en.wikipedia.org/wiki/Discounted_cumulative_gain
        :param relevance:
        :param p:
        """
        self.relevance = np.asfarray(relevance)
        self.p = p

    def CG(self):
        """
        Cumulative Gain
        :return:
        """
        return np.sum(self.relevance[:self.p])

    def DCG(self):
        """
        Discounted Cumulative Gain
        :return:
        """
        positions = np.arange(1, 1 + self.p, 1, np.float)
        relevances = self.relevance[:self.p]
        return np.sum(relevances / np.log2(positions + 1))

    def IDCG(self):
        """
        Ideal Discounted Cumulative Gain
        :return:
        """
        positions = np.arange(1, 1 + self.p, 1, np.float)
        relevances = self.relevance[np.argsort(-self.relevance)][:self.p]
        return np.sum(relevances / np.log2(positions + 1))

    def NDCG(self):
        """
        Normalized Discounted Cumulative Gain
        :return:
        """
        return self.DCG() / self.IDCG()


if __name__ == '__main__':
    solution = Solution([3, 2, 3, 0, 1, 2, 3, 2])
    print(solution.CG())
    print(solution.DCG())
    print(solution.IDCG())
    print(solution.NDCG())