机器学习：决策树算法代码详细注释笔记

计算给定数据集的香农熵

# 计算给定数据集的香农熵
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 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]
        # 属性值去重
        uniqueVals = set(featList)
        # 以当前属性分类数据集，计算信息熵
        newEntropy = 0.0
        # 遍历当前属性的每一个值计算信息熵
        for value in uniqueVals:
            # 数据集中属性i值等于当前value的数据集
            subDataSet = splitDataSet(dataSet, i, value)
            # 计算当前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]

创建决策树递归函数

# 创建决策树递归函数
# dataSet 用于分类训练的数据集
# labels 每一列的特征名称数组 (为了使数据具有比较明确的含义)
def createTree(dataSet, labels):
    # 取出数据集最后一列
    classList = [example[-1] for example in dataSet]
    # 递归函数的第一个停止条件: 当前数据集的所有标签类完全相同，则没有必要继续分类下去
    if (classList.count(classList[0]) == len(classList)):
        return classList[0]
    # 递归函数的第二个停止条件: 当使用完了所有特征，仍不能将数据集划分成包含唯一类别的分组
    # len(dataSet[0]) == 1 表示数据集的属性只剩下一个无法继续分类
    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

分类算法

# 分类算法
# inputTree 训练好的决策树
# featLabels 特征项数组
# testVec 测试向量
def classify(inputTree, featLabels, testVec):
    # 获取树根节点名称
    firstStr = list(inputTree.keys())[0]
    # 获取第二级子节点
    secondDict = inputTree[firstStr]
    # 获取特征项下标
    featIndex = featLabels.index(firstStr)
    # 获取当前树第一级特征项的所有分类值
    for key in secondDict.keys():
        # 如果特征向量的属性值 等于当前值
        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):
    import pickle
    fw = open(filename, 'wb')
    pickle.dump(inputTree, fw)
    fw.close()

加载决策树

# 加载训练好的决策树
def garbTree(filename):
    import pickle 
    fr = open(filename, 'rb+')
    return pickle.load(fr)

未完待续...