1. from math import log
    2. import operator
    5. def createDataSet():
    6.     dataSet = [[0, 0, 0, 0, 'N'],
    7.                [0, 0, 0, 1, 'N'],
    8.                [1, 0, 0, 0, 'Y'],
    9.                [2, 1, 0, 0, 'Y'],
    10.                [2, 2, 1, 0, 'Y'],
    11.                [2, 2, 1, 1, 'N'],
    12.                [1, 2, 1, 1, 'Y']]
    13.     labels = ['outlook', 'temperature', 'humidity', 'windy']
    14.     return dataSet, labels
    17. def calcShannonEnt(dataSet):  # 计算熵
    18.     numEntries = len(dataSet)
    19.     labelCounts = {}
    20.     for featVec in dataSet:
    21.         currentLabel = featVec[-1]
    22.         if currentLabel not in labelCounts.keys():
    23.             labelCounts[currentLabel] = 0
    24.         labelCounts[currentLabel] += 1  # 数每一类各多少个， {'Y': 4, 'N': 3}
    25.     shannonEnt = 0.0
    26.     for key in labelCounts:
    27.         prob = float(labelCounts[key]) / numEntries
    28.         shannonEnt -= prob * log(prob, 2)
    29.     return shannonEnt
    32. def chooseBestFeatureToSplit(dataSet):
    33.     numFeatures = len(dataSet[0]) - 1  # feature个数
    34.     baseEntropy = calcShannonEnt(dataSet)  # 整个dataset的熵
    35.     bestInfoGainRatio = 0.0
    36.     bestFeature = -1
    37.     for i in range(numFeatures):
    38.         featList = [example[i] for example in dataSet]  # 每个feature的list
    39.         uniqueVals = set(featList)  # 每个list的唯一值集合
    40.         newEntropy = 0.0
    41.         splitInfo = 0.0
    42.         for value in uniqueVals:
    43.             subDataSet = splitDataSet(dataSet, i, value)  # 每个唯一值对应的剩余feature的组成子集
    44.             prob = len(subDataSet) / float(len(dataSet))
    45.             newEntropy += prob * calcShannonEnt(subDataSet)
    46.             splitInfo += -prob * log(prob, 2)
    47.         infoGain = baseEntropy - newEntropy  # 这个feature的infoGain
    48.         if (splitInfo == 0):  # fix the overflow bug
    49.             continue
    50.         infoGainRatio = infoGain / splitInfo  # 这个feature的infoGainRatio增益率
    51.         if (infoGainRatio > bestInfoGainRatio):  # 选择最大的gain ratio
    52.             bestInfoGainRatio = infoGainRatio
    53.             bestFeature = i  # 选择最大的gain ratio对应的feature
    54.     return bestFeature
    57. def splitDataSet(dataSet, axis, value):
    58.     retDataSet = []
    59.     for featVec in dataSet:
    60.         if featVec[axis] == value:  # 只看当第i列的值＝value时的item
    61.             reduceFeatVec = featVec[:axis]  # featVec的第i列给除去
    62.             reduceFeatVec.extend(featVec[axis + 1:])
    63.             retDataSet.append(reduceFeatVec)
    64.     return retDataSet
    67. def createTree(dataSet, labels):
    68.     classList = [example[-1] for example in dataSet]  # ['N', 'N', 'Y', 'Y', 'Y', 'N', 'Y']
    69.     if classList.count(classList[0]) == len(classList):
    70.         # classList所有元素都相等，即类别完全相同，停止划分
    71.         return classList[0]  # splitDataSet(dataSet, 0, 0)此时全是N，返回N
    72.     if len(dataSet[0]) == 1:  # [0, 0, 0, 0, 'N']
    73.         # 遍历完所有特征时返回出现次数最多的
    74.         return majorityCnt(classList)
    75.     bestFeat = chooseBestFeatureToSplit(dataSet)  # 0－> 2
    76.     # 选择最大的gain ratio对应的feature
    77.     bestFeatLabel = labels[bestFeat]  # outlook -> windy
    78.     myTree = {bestFeatLabel: {}}
    79.     # 多重字典构建树{'outlook': {0: 'N'
    80.     del (labels[bestFeat])  # ['temperature', 'humidity', 'windy'] -> ['temperature', 'humidity']
    81.     featValues = [example[bestFeat] for example in dataSet]  # [0, 0, 1, 2, 2, 2, 1]
    82.     uniqueVals = set(featValues)
    83.     for value in uniqueVals:
    84.         subLabels = labels[:]  # ['temperature', 'humidity', 'windy']
    85.         myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels)
    86.         # 划分数据，为下一层计算准备
    87.     return myTree
    90. def majorityCnt(classList):  # 如果属性完全相同，却不具有相同的类别，则采用少数服从多数的原则进行划分
    91.     classCount = {}
    92.     for vote in classList:
    93.         if vote not in classCount.keys():
    94.             classCount[vote] = 0
    95.         else:
    96.             classCount[vote] += 1
    97.     sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    98.     return sortedClassCount[0][0]

    https://github.com/cdqncn/JueCeShu/blob/master/myTree.py