python实现香农熵计算

超、凢脫俗 2022-02-24 15:40 307阅读 0赞

#!/usr/bin/env python
    # encoding: utf-8
    """
    @author: wanwei
    @license: (C) Copyright 2013-2017, Node Supply Chain Manager Corporation Limited.
    @software: pycharm
    @file: shangnon.py
    @time: 2019/4/4 14:03
    @desc:
    """
    import math
    
    def createDataSet():
        """
        创建数据集
        :return:
        """
        dataSet = [[u'青年',u'否',u'否',u'一般',u'拒绝'],
                   [u'青年', u'否', u'否', u'好', u'同意'],
                   [u'青年', u'是', u'否', u'好', u'拒绝'],
                   [u'青年', u'是', u'是', u'好', u'同意'],
                   [u'青年', u'是', u'是', u'一般', u'同意'],
                   [u'青年', u'是', u'否', u'一般', u'同意'],
                   [u'中年', u'否', u'否', u'一般', u'同意'],
                   [u'中年', u'否', u'否', u'好', u'拒绝'],
                   [u'中年', u'是', u'否', u'好', u'拒绝'],
                   [u'中年', u'是', u'是', u'好', u'拒绝'],
                   [u'中年', u'是', u'是', u'一般', u'拒绝'],
                   [u'中年', u'是', u'否', u'一般', u'拒绝'],
                   [u'老年', u'否', u'否', u'一般', u'同意'],
                   [u'老年', u'否', u'否', u'好', u'拒绝'],
                   [u'老年', u'是', u'否', u'好', u'拒绝'],
                   [u'老年', u'是', u'是', u'好', u'拒绝'],
                   [u'老年', u'是', u'是', u'一般', u'同意'],
                   [u'老年', u'是', u'否', u'一般', u'拒绝'],
                   ]
        labelSet = [u'年龄',u'有工作',u'有房子',u'信贷情况']
        # 返回数据集和标签
        return dataSet, labelSet
    
    def cacShannonEnt(dataSet):
        """
        计算训练数据集Y随机变量的香农熵
        :param dataSet:
        :return:
        """
        numEntries = len(dataSet)
        labelCounts = {}
        for featvec in dataSet:
            currentLbel = featvec[-1]
            if currentLbel not in labelCounts.keys():
                labelCounts[currentLbel] = 0
            labelCounts[currentLbel] += 1
        shannonEnt = 0.0
        for key in labelCounts:
            prob = float(labelCounts[key])/numEntries
            shannonEnt -= prob * math.log(prob, 2) # log base 2
        return shannonEnt
    
    if __name__ == '__main__':
        dataSet, labelSet = createDataSet()
        print(cacShannonEnt(dataSet))