'''
优点：计算复杂度不高，输出结果易于理解，对中间值的缺失不敏感，可以处理不相关特征数据
缺点：可能会产生过度匹配问题
适用数据类型：数值型和标称型
信息增益：ID3   
信息增率：C4.5
基尼指数：CART
'''
from math import log
import operator

#计算给定数据集的香农熵
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)  #以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

def splitDataSet(dataSet, axis, value): 
    #输入参数：待划分的数据集，划分数据集的特征，特征的返回值
    retDataSet = [] #防止原始数据集被修改，声明一个新的列表对象
    for featVec in dataSet:  #遍历数据集中的每个元素
        if featVec[axis] == value:  #抽取符合特征的数据
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis + 1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet

def chooseBestFeatureToSplit(dataSet):
    #选择最好的数据集划分方式
   '''
   调用时数据需要满足的要求：
   数据必须是由相同数据长度的列表元素组成的列表，
   数据的最后一列或每个实例的最后一个元素是类别标签
   '''
   numFeatures = len(dataSet[0]) - 1
   baseEntropy = calcShannonEnt(dataSet)  #计算整个数据集的香农熵
   bestInfoGain = 0.0
   bestFeature = -1
   for i in range(numFeatures): #遍历数据集的所有特征
        featList = [example[i] for example in dataSet] #将数据集中所有第i个特征值写入新列表中
        uniqueVals = set(featList)  #将列表转为集合数据类型，快速得到列表中唯一元素值
        newEntropy = 0.0
        for value in uniqueVals:  #遍历当前特征中所有的唯一属性值
            subDataSet = splitDataSet(dataSet, i, value)  #对每个特征划分一次数据集
            prob = len(subDataSet)/float(len(dataSet))
            newEntropy += prob * calcShannonEnt(subDataSet) #计算数据集的新熵值，并对所有唯一特征值得到的熵求和
        infoGain = baseEntropy - newEntropy  #计算信息增益，熵的减少
        if (infoGain > bestInfoGain):  #比较所有特征的信息增益
            bestInfoGain = infoGain
            bestFeature = i
   return bestFeature

def majorityCnt(classList):
    #相当于投票表决，采用多数表决的方法决定该叶节点的分类
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys():classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.items(), key = operator.itemgetter(1), reverse = True)
    return sortedClassCount[0][0]

def createTree(dataSet,labels):   
#创建树，输入参数为数据集和标签列表
    classList = [example[-1] for example in dataSet] #创建包含数据集所有类标签的列表
    if classList.count(classList[0]) == len(classList):  #类别完全相同则停止继续划分
        return classList[0]
    if len(dataSet[0]) == 1:   #遍历完所有特征时返回出现次数最多的标签
        return majorityCnt(classList)
    bestFeat = chooseBestFeatureToSplit(dataSet)  #选取最好的特征划分
    bestFeatLabel = labels[bestFeat]  #最好特征对应的标签
    myTree = {bestFeatLabel:{}}  #用于存储树的信息的字典变量
    del(labels[bestFeat])
    featValues = [example[bestFeat] for example in dataSet] 
    uniqueVals = set(featValues)  #得到列表包含的所有属性值
    for value in uniqueVals:
        subLabels = labels[:] #为了不改变原始列表内容，复制类标签，使用新变量代替原始列表
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels)
    return myTree
   
def classify(inputTree, featLabels, testVec):
    firstSides = list(inputTree.keys())
    firstStr = firstSides[0]  #第一次划分数据集的类标签
    secondDict = inputTree[firstStr]  #第一次划分数据集的类标签所附带的子节点取值
    featIndex = featLabels.index(firstStr)  #使用index方法查找当前列表中第一个匹配firstStr变量的元素
    for key in secondDict.keys():  #递归遍历整棵树，比较testVec变量中的值与树节点的值
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ == 'dict':
                classLabel = classify(secondDict[key], featLabels, testVec)
            else:   classLabel = secondDict[key]  #如果到达叶子节点，则返回当前节点的分类标签
    return classLabel

def storeTree(inputTree, filename):
#使用模块pickle序列化对象，序列化对象可以在磁盘上保存对象并在需要的时候读取出来    
    import pickle  
    fw = open(filename,'w')
    pickle.dump(inputTree, fw)
    fw.close()
    
def grabTree(filename):
    import pickle
    fr = open(filename)
    return pickle.load(fr)

决策树可视化

import matplotlib.pyplot as plt

decisionNode = dict(boxstyle = 'sawtooth', fc = '0.8')  #定义文本框和箭头格式
leafNode = dict(boxstyle = 'round4', fc = '0.8')
arrow_args = dict(arrowstyle = '<-')

def plotNode(nodeTxt, centerPt, parentPt, nodeType):  #绘制带箭头的注解
    createPlot.ax1.annotate(nodeTxt, xy = parentPt, xycoords = 'axes fraction', 
                            xytext = centerPt, textcoords = 'axes fraction',
                            va = 'center', ha = 'center', bbox = nodeType,arrowprops = arrow_args)
 
def getNumLeafs(myTree):
#获取叶节点的数目
    numLeafs = 0
    firstSides = list(myTree.keys())
    firstStr = firstSides[0]  #第一次划分数据集的类标签
    secondDict = myTree[firstStr]   #第一次划分数据集的类标签所附带的子节点取值
    for key in secondDict.keys(): #遍历整棵树的所有子节点
        if type(secondDict[key]).__name__== 'dict':  #如果子节点是字典类型，则该节点是一个判断结点，否则为叶节点
            numLeafs += getNumLeafs(secondDict[key])
        else:   numLeafs += 1  #累计叶子节点的个数并返回该数值
    return numLeafs

def getTreeDepth(myTree):
#获取树的层数
    maxDepth = 0
    firstSides = list(myTree.keys())
    firstStr = firstSides[0]  #第一次划分数据集的类标签
    secondDict = myTree[firstStr]  #第一次划分数据集的类标签所附带的子节点取值
    for key in secondDict.keys():  #统计遍历过程中遇到判断节点的个数
        if type(secondDict[key]).__name__== 'dict':
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:   thisDepth = 1
        if thisDepth > maxDepth: maxDepth = thisDepth
    return maxDepth
    
def retrieveTree(i):
#预先存储的树的信息
    listOfTrees = [{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},
                   {'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}]
    return listOfTrees[i]

def plotMidText(cntrPt, parentPt, txtString):
#计算父节点和子节点的中间位置，并在此处添加简单的文本标签信息
    xMid = (parentPt[0] - cntrPt[0])/2.0 + cntrPt[0]
    yMid = (parentPt[1] - cntrPt[1])/2.0 + cntrPt[1]
    createPlot.ax1.text(xMid, yMid, txtString)
    
def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)   #计算树的宽度
    depth = getTreeDepth(myTree)  #计算树的深度
    firstSides = list(myTree.keys())
    firstStr = firstSides[0]  #第一次划分数据集的类标签
    #追踪已经绘制的节点位置，以及放置下一个节点的恰当位置
    cntrPt = (plotTree.xOff + (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yOff)
    #树的宽度用于计算放置判断节点的位置，原则是将它放在所有叶子节点的中间
    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff - 1.0/plotTree.totalD  #由于是自顶向下绘制图形，故依次递减y坐标
    for key in secondDict.keys():
        if type(secondDict[key]).__name__=='dict': #如果节点是判断节点则递归调用plotTree()函数
            plotTree(secondDict[key], cntrPt, str(key))
        else:  #如果节点是叶子节点，则在图形上画出叶子节点
            plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW
            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
    plotTree.yOff = plotTree.yOff + 1.0/plotTree.totalD  #绘制了所有子节点后，增加y坐标的偏移
   
def createPlot(inTree):
    fig = plt.figure(1,facecolor = 'white')
    fig.clf()  #创建新图形并清空绘图区
    axprops = dict(xticks = [], yticks = [])
    createPlot.ax1 = plt.subplot(111, frameon = False, **axprops)
    #使用树的宽度与深度计算树节点的摆放位置
    plotTree.totalW = float(getNumLeafs(inTree))  #存储树的宽度
    plotTree.totalD = float(getTreeDepth(inTree))  #存储树的深度
    plotTree.xOff = -0.5/plotTree.totalW
    plotTree.yOff = 1.0
    plotTree(inTree,(0.5,1.0), '')
    plt.show()