Calculate normalized Discounted Cumulative Gain (nDCG)

ndcg.py

import math

user1 = [1,1,1,0,0,1,0]
user2 = [1,1,0,0,0,0,0]
user3 = [0,0,1,1,0,1,0]

def calcNdcg(user):
    idcg = getIdcg(len(user))
    dcg = getDcg(user)
    ndcg = dcg / idcg
    print(ndcg)
    return ndcg

def getIdcg(n):
    idcg = 0
    for i in range(1, n + 1):
        idcg = idcg + 1 / math.log2(i + 1)
    return idcg
    
def getDcg(user):
    dcg = 0
    for i in range(1, len(user) + 1):
        dcg = dcg + user[i - 1] / math.log2(i + 1)
    return dcg



n1 = calcNdcg(user1)
n2 = calcNdcg(user2)
n3 = calcNdcg(user3)
print("Average: " + str((n1 + n2 + n3) / 3))