计算KL散度

计算KL散度

import math
p=list(map(int,input().strip().split()))
q=list(map(int,input().strip().split()))

def cal_prob(ss):
    s_prob={}
    for item in ss:
        if item not in s_prob.keys():
            s_prob[item] = 1
        else:
            s_prob[item] += 1
    for key in s_prob.keys():
        s_prob[key]=s_prob[key]/len(ss)

    return s_prob

def cal_kl(p,q):
    res=0
    for key in p.keys():
        if key in q.keys():
            res+=p[key]*math.log(p[key]/q[key],2)
    return res

pp = cal_prob(p)
qp = cal_prob(q)
print(round(cal_kl(pp,qp),2))

测试用例：

1 1 1 2 3 4 1 2 3 4 1 3 4 1 2 3 3 1 2 3 1 2 3 1 3 1
2 2 2 2 3 4 4 2 2 1 3 4 2 2 2 1 3 4 4 2 4 3