1. import pandas as pd
    2. from math import log
    3. import treePlotter
    4. import operator
    6. #计算给定数据集的香农熵
    7. def calcShannonEnt(dataSet):
    8.     numEntries = len(dataSet)
    9.     labelCounts = {}
    10.     for featVec in dataSet:
    11.         currentLabel = featVec[-1]
    12.         if currentLabel not in labelCounts.keys():
    13.             labelCounts[currentLabel] = 0
    14.         labelCounts[currentLabel] += 1
    15.     shannonEnt = 0.0
    16.     for key in labelCounts:
    17.         prob = float(labelCounts[key])/numEntries
    18.         shannonEnt -= prob * log(prob, 2)
    19.     return shannonEnt
    21. # 对离散变量划分数据集，取出该特征取值为value的所有样本
    22. def splitDataSet(dataSet, axis, value):
    23.     retDataSet = []
    24.     for featVec in dataSet:
    25.         if featVec[axis] == value:
    26.             reducedFeatVec = featVec[:axis]
    27.             reducedFeatVec.extend(featVec[axis + 1:])
    28.             retDataSet.append(reducedFeatVec)
    29.     return retDataSet
    31. # 对连续变量划分数据集——二分法。不大于或者大于value的样本分别保存，进行划分
    32. # direction规定划分的方向，决定是划分出小于value的数据样本还是大于value的数据样本集
    33. def splitContinuousDataSet(dataSet, axis, value, direction):
    34.     retDataSet = []
    35.     for featVec in dataSet:
    36.         if direction == 0:
    37.             if featVec[axis] > value:
    38.                 retDataSet.append(featVec)
    39.         else:
    40.             if featVec[axis] <= value:
    41.                 retDataSet.append(featVec)
    42.     return retDataSet
    44. # 选择最好的数据集划分方式
    45. def chooseBestFeatureToSplit(dataSet, labels):
    46.     numFeatures = len(dataSet[0]) - 1
    47.     baseEntropy = calcShannonEnt(dataSet)
    48.     bestInfoGain = 0.0
    49.     bestFeature = -1
    50.     bestSplitDict = {}
    51.     for i in range(numFeatures):
    52.         featList = [example[i] for example in dataSet]
    53.         # 对连续型特征进行处理
    54.         if type(featList[0]).__name__ == 'float' or type(featList[0]).__name__ == 'int':
    55.             sortfeatList = sorted(featList)  # 二分法：先对属性值从小到大进行排序
    56.             splitList = []
    57.             for j in range(len(sortfeatList) - 1):
    58.                 splitList.append((sortfeatList[j] + sortfeatList[j + 1]) / 2.0)
    59.             bestSplitEntropy = 10000
    60.             slen = len(splitList)
    61.             for j in range(slen):
    62.                 value = splitList[j]
    63.                 newEntropy = 0.0
    64.                 subDataSet0 = splitContinuousDataSet(dataSet, i, value, 0)
    65.                 subDataSet1 = splitContinuousDataSet(dataSet, i, value, 1)
    66.                 prob0 = len(subDataSet0) / float(len(dataSet))
    67.                 prob1 = len(subDataSet1) / float(len(dataSet))
    68.                 newEntropy = prob0 * calcShannonEnt(subDataSet0) \
    69.                     + prob1 * calcShannonEnt(subDataSet1)
    70.                 if newEntropy < bestSplitEntropy:
    71.                     bestSplitEntropy = newEntropy
    72.                     bestSplit = j
    73.             bestSplitDict[labels[i]] = splitList[bestSplit]
    74.             infoGain = baseEntropy - bestSplitEntropy  # 当前连续型特征最优划分点的信息增益
    75.         # 对离散型特征进行处理
    76.         else:
    77.             uniqueVals = set(featList)
    78.             newEntropy = 0.0
    79.             for value in uniqueVals:
    80.                 subDataSet = splitDataSet(dataSet, i, value)
    81.                 prob = len(subDataSet) / float(len(dataSet))
    82.                 newEntropy += prob * calcShannonEnt(subDataSet)
    83.             infoGain = baseEntropy - newEntropy
    84.         #  比较每一个特征的信息增益，选择最大的。
    85.         if infoGain > bestInfoGain:
    86.             bestInfoGain = infoGain
    87.             bestFeature = i
    88.     # 返回：最优划分特征所在的列，和连续型{特征标签:最优划分点}
    89.     return bestFeature, bestSplitDict
    91. # 采用多数表决（投票）的方法决定该叶子结点的分类。
    92. def majorityCnt(classList):
    93.     classCount = {}
    94.     for vote in classList:
    95.         if vote not in classCount.keys():
    96.             classCount[vote] = 0
    97.         classCount[vote] += 1
    98.     sortedClassCount = sorted(classCount.items(), \
    99.         key=operator.itemgetter(1), reverse=True)
    100.     return sortedClassCount[0][0]
    102. # 主程序，递归产生决策树
    103. def createTree(dataSet, labels, data_full, labels_full):
    104.     classList = [example[-1] for example in dataSet]
    105.     if classList.count(classList[0]) == len(classList):
    106.         # 一、如果样本都属于同一类，就没必要再分类了。
    107.         return classList[0]
    108.     if len(dataSet[0]) == 1:
    109.         # 二、所有特征都用完了，但类别标签仍然不是唯一的,返回出现次数最多的类别
    110.         return majorityCnt(classList)
    111.     bestFeat, bestSplitDict = chooseBestFeatureToSplit(dataSet, labels)
    112.     if bestFeat == -1:
    113.         # 三、没有选出最优划分特征，返回出现次数最多的类别
    114.         return majorityCnt(classList)
    115.     bestFeatLabel = labels[bestFeat]
    116.     # 1、最优划分特征是离散型
    117.     if type(dataSet[0][bestFeat]).__name__ == 'str':
    118.         myTree = {bestFeatLabel: {}}
    119.         featValues = [example[bestFeat] for example in dataSet]
    120.         uniqueVals = set(featValues)
    121.         bestFeatIndexInFull = labels_full.index(bestFeatLabel)
    122.         featValuesFull = [example[bestFeatIndexInFull] \
    123.             for example in data_full]
    124.         uniqueValsFull = set(featValuesFull)
    125.         del (labels[bestFeat])
    126.         for value in uniqueValsFull:
    127.             if value in uniqueVals:
    128.                 subLabels = labels[:]
    129.                 valueDataSet = splitDataSet(dataSet, bestFeat, value)
    130.                 myTree[bestFeatLabel][value] = createTree(valueDataSet, \
    131.                     subLabels, data_full, labels_full)
    132.             else:
    133.                 myTree[bestFeatLabel][value] = majorityCnt(classList)
    134.     # 2、连续型特征，不删除
    135.     else:
    136.         bestSplitValue = bestSplitDict[bestFeatLabel]
    137.         bestFeatLabel = labels[bestFeat] + '<=' + str(bestSplitValue)
    138.         myTree = {bestFeatLabel: {}}
    139.         subDataSet0 = splitContinuousDataSet(dataSet, bestFeat, bestSplitValue, 0)  # >value
    140.         subDataSet1 = splitContinuousDataSet(dataSet, bestFeat, bestSplitValue, 1)  # <=value
    141.         myTree[bestFeatLabel]['no'] = createTree(subDataSet0, labels, \
    142.             data_full, labels_full)
    143.         myTree[bestFeatLabel]['yes'] = createTree(subDataSet1, labels, \
    144.             data_full, labels_full)
    145.     return myTree
    147. #决策树
    148. def classify(inputTree, featLabels, testVec):
    149.     firstStr = list(inputTree.keys())[0]
    150.     secondDict = inputTree[firstStr]
    151.     featIndex = featLabels.index(firstStr)
    152.     for key in secondDict.keys():
    153.         if testVec[featIndex] == key:
    154.             if type(secondDict[key]).__name__ == 'dict':
    155.                 classLabel = classify(secondDict[key], featLabels, testVec)
    156.             else:
    157.                 classLabel = secondDict[key]
    158.     return classLabel
    160. #决策树
    161. def classifyAll(inputTree, featLabels, testDataSet):
    162.     classLabelAll = []
    163.     for testVec in testDataSet:
    164.         classLabelAll.append(classify(inputTree, featLabels, testVec))
    165.     return classLabelAll
    167. if __name__ == '__main__':
    168.     file = pd.read_csv('data.csv')
    169.     data = file.values[:, 3:].tolist()
    170.     for i in [t[0:-1] for t in data]: #连续型变量转化为浮点型
    171.         for j in i:
    172.             j = float(j)
    173.     label = file.columns.values[3:].tolist()
    174.     label_full = label
    175.     train_data = data[1:200]
    176.     train_data_full = train_data
    177.     test_data = data[200:245]
    178.     desicionTree = createTree(train_data, label, train_data_full, label_full)
    179.     print('desicionTree:\n', desicionTree)
    180.     treePlotter.createPlot(desicionTree)
    181.     print('classifyResult:\n', classifyAll(desicionTree, label, test_data))