决策树算法学习笔记（一）

香农熵的计算

# -*- coding: UTF-8 -*-
from math import log

def calcShannonEnt(dataSet):
    numEntries=len(dataSet)#统计数据集的数量
    labelCounts={}#创建一个数据字典
    for featVec in dataSet:
        currentLabel=featVec[-1]#取数据集的最后一列数据
        if currentLabel not in labelCounts.keys():#数据字典中是否存在当前键值
            labelCounts[currentLabel]=0#数据字典中如果不存在则加入字典中
        labelCounts[currentLabel]+=1#数据字典中如果存在当前键值则数量加一
    shannonEnt=0.0#初始化香农熵
    for key in labelCounts:#计算香农熵
        prob=float(labelCounts[key])/numEntries
        shannonEnt-=prob*log(prob,2)
    return shannonEnt

def createDataSet():
    dataSet=[[1,1,'yes'],[1,1,'yes'],[1,0,'no'],[0,1,'no'],[0,1,'no']]
    labels=['no surfacing','flippers']
    return dataSet,labels
myDat,labels=createDataSet()
print calcShannonEnt(myDat)
0.970950594455