python ID3 决策树 代码

 代码参考了python实现ID3决策树分类算法_aoanng的博客-CSDN博客_id3算法python实现

有所精简 

'''
function：ID3决策树生成算法
author：baomi
date: 2021/11/01
reference: https://blog.csdn.net/colourful_sky/article/details/82056125
'''

import math

def splitDataSet(dataSet, i, value):
    '''
    返回数据集dataSet中，去掉第i列属性值为value的实例后形成的新的数据集
    '''
    retDataSet = []
    for x in dataSet:
        if x[i] == value:
            temp = x[:]
            temp.pop(i)
            retDataSet.append(temp)
    return retDataSet


def calcEntropy(dataSet):
    '''
    计算一个数据集的熵
    '''
    labelDict = {}  # 数据集的标签-该标签总个数
    for x in dataSet:
        label = x[-1]
        if label not in labelDict.keys():
            labelDict[label] = 0
        labelDict[label] += 1

    n = len(dataSet)
    retEntropy = 0.0
    for key in labelDict:
        p = float(labelDict[key]) / n  # 计算标签概率
        retEntropy -= p * math.log(p, 2)

    return retEntropy


def calcInfoGain(dataSet, i):
    '''
    计算对数据集dataSet，选定第i列特征时所获得的信息增益
    '''
    preEntropy = calcEntropy(dataSet)

    postEntropy = 0.0
    featureSet = set([x[i] for x in dataSet])  # 得到i列特征所有特征值的集合
    for feature in featureSet:  # 以feature为筛选条件，计算筛选后的数据集熵
        subDataSet = splitDataSet(dataSet, i, feature)
        subDataSetEntropy = calcEntropy(subDataSet)
        p = len(subDataSet)/len(dataSet)
        postEntropy += p*subDataSetEntropy

    return preEntropy-postEntropy

def getMaxInfoGainNode(dataSet, featureNameList):
    '''
    featureNameList是dataSet中各特征名称
    该函数返回两种结果：
    熵为0时，返回标签,类型为str
    熵不为0时，返回具有最大信息增益的特征的索引号，特征名，以及最大信息增益
    '''
    dataSetEntropy = calcEntropy(dataSet)
    if dataSetEntropy == 0:
        return dataSet[0][-1]  # 数据集熵为0，说明标签都相同，直接将该标签返回

    featureNum = len(featureNameList)

    maxInfoGain = 0
    maxInfoGainIndex = 0
    for i in range(0, featureNum): #遍历所有特征，获得具有最大信息增益的特征索引号
        infoGain = calcInfoGain(dataSet, i)
        if infoGain > maxInfoGain:
            maxInfoGain = infoGain
            maxInfoGainIndex = i

    return maxInfoGainIndex, featureNameList[maxInfoGainIndex], maxInfoGain


def createID3Tree(dataSet, featureNameList):
    '''
    该函数返回一个结点
    如果dataSet熵为0，那么返回dataSet中类标签，此标签唯一
    否则，返回一个字典，该字典的key为dataSet选出的最优特征名，value又为一个字典，
    value字典的key为最优特征的特征值名，value字典的value又为一个字典.....
    '''
    maxInfoGainNode = getMaxInfoGainNode(dataSet, featureNameList)
    if type(maxInfoGainNode) == str:
        return maxInfoGainNode

    nodeIndex, nodeName = maxInfoGainNode[0], maxInfoGainNode[1]
    ret = {}
    ret[nodeName] = {}

    featureSet = set([x[nodeIndex] for x in dataSet])
    for feature in featureSet:
        subDataSet = splitDataSet(dataSet, nodeIndex, feature)
        newFeatNameList = featureNameList[:]
        newFeatNameList.pop(nodeIndex)
        childTree = createID3Tree(subDataSet, newFeatNameList) #对以最大信息增益作为特征筛选后的子数据集进行递归调用
        ret[nodeName][feature] = childTree

    return ret

dataSet = [['青年', '否', '否', '一般', '拒绝'],
           ['青年', '否', '否', '好', '拒绝'],
           ['青年', '是', '否', '好', '同意'],
           ['青年', '是', '是', '一般', '同意'],
           ['青年', '否', '否', '一般', '拒绝'],
           ['中年', '否', '否', '一般', '拒绝'],
           ['中年', '否', '否', '好', '拒绝'],
           ['中年', '是', '是', '好', '同意'],
           ['中年', '否', '是', '非常好', '同意'],
           ['中年', '否', '是', '非常好', '同意'],
           ['老年', '否', '是', '非常好', '同意'],
           ['老年', '否', '是', '好', '同意'],
           ['老年', '是', '否', '好', '同意'],
           ['老年', '是', '否', '非常好', '同意'],
           ['老年', '否', '否', '一般', '拒绝'], ]

featureNameList = ['年龄', '有工作', '有房子', '信贷情况']


ID3Tree = createID3Tree(dataSet, featureNameList)
print(ID3Tree)

输出结果：

对结果用matplotlib绘图

代码来自Matplotlib绘制树形图_wancongconghao的博客-CSDN博客_matplotlib 树状图 

#绘制树形图

import matplotlib.pyplot as plt

decision_node = dict(boxstyle="sawtooth",fc="0.8")

leaf_node = dict(boxstyle="round4",fc="0.8")

arrow_args = dict(arrowstyle="<-")

#获取树的叶子结点个数（确定图的宽度）

def get_leaf_num(tree):

    leaf_num = 0

    first_key = list(tree.keys())[0]

    next_dict = tree[first_key]

    for key in next_dict.keys():

        if type(next_dict[key]).__name__=="dict":

            leaf_num +=get_leaf_num(next_dict[key])

        else:

            leaf_num +=1

    return leaf_num

#获取数的深度（确定图的高度）

def get_tree_depth(tree):

    depth = 0

    first_key = list(tree.keys())[0]

    next_dict = tree[first_key]

    for key in next_dict.keys():

        if type(next_dict[key]).__name__ == "dict":

            thisdepth = 1+ get_tree_depth(next_dict[key])

        else:

            thisdepth = 1

        if thisdepth>depth: depth = thisdepth

    return depth

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 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, va="center", ha="center", rotation=30)

def plotTree(myTree, parentPt, nodeTxt):

    numLeafs = get_leaf_num(myTree)

    depth = get_tree_depth(myTree)

    firstStr = list(myTree.keys())[0]

    cntrPt = (plotTree.xOff + (1.0 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yOff)

    plotMidText(cntrPt, parentPt, nodeTxt)

    plotNode(firstStr, cntrPt, parentPt, decision_node)

    secondDict = myTree[firstStr]

    plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD

    for key in secondDict.keys():

        if type(secondDict[

                    key]).__name__ == 'dict':

            plotTree(secondDict[key], cntrPt, str(key))

        else:

            plotTree.xOff = plotTree.xOff + 1.0 / plotTree.totalW

            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leaf_node)

            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))

    plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD

 

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(get_leaf_num(inTree))

    plotTree.totalD = float(get_tree_depth(inTree))

    plotTree.xOff = -0.5 / plotTree.totalW

    plotTree.yOff = 1.0

    plotTree(inTree, (0.5, 1.0), '')

    plt.show()

plt.rcParams['font.sans-serif']=['SimHei'] #用来正常显示中文标签

plt.rcParams['axes.unicode_minus']=False #用来正常显示负号

tree = {'有房子': {'否': {'有工作': {'否': '拒绝', '是': '同意'}}, '是': '同意'}}

createPlot(tree)

结果：