Decision Tree 决策树算法Python实现

决策树是很自然的逻辑结构，本文代码是Machine Learning in Action 这本书中算法的实现

另作补充：每次最佳划分节点的确定方式（来自周志华《机器学习》）：

通过最大化信息增益来选择划分特征

#!/usr/bin/python
# -*- coding: utf-8 -*-
from math import log
#from numpy import*
import operator
import matplotlib
import matplotlib.pyplot as plt
from os import listdir
def calcShannonEnt(dataSet):   #计算香农熵
    numEntries = len(dataSet)
    labelCounts = {}
    for featVec in dataSet: #the the number of unique elements and their occurance
        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) #log base 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):   #选择一个最优的特征划分，对于一个特征i，枚举所有该特征下的取值，累加每个取值下数据的熵。
    numFeatures = len(dataSet[0]) - 1
    baseEntropy = calcShannonEnt(dataSet)
    bestInfoGain = 0.0; bestFeature = -1
    for i in range(numFeatures):  #遍历特征i
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)
        newEntropy = 0.0
        for value in uniqueVals:   #遍历特征i的所有取值
            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.iteritems(), 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):     #用已有的决策树分类
    firstStr=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,'w')
    pickle.dump(inputTree,fw)
    fw.close()
def grabTree(fliename):
    import pickle
    fr=open(fliename)
    return pickle.load(fr)

myDat,labels=createDataSet()


#print myDat ,myDat[0]
# # print calcShannonEnt(myDat)
# print splitDataSet(myDat,0,1)
# print splitDataSet(myDat,0,0)
# aa=array(myDat)
# print aa[0]
# print len(aa[0])
nn=labels[:]
mytree=createTree(myDat,nn)
storeTree(mytree,'aa.txt')
print grabTree('aa.txt')
print mytree,labels
print classify(mytree,labels,[1,0])
print classify(mytree,labels,[1,1])