决策树, Splitting Datasets one feature at a time

'''
决策树算法实现
'''
from numpy import *
from math import log
import operator


# 计算香农熵
def calcShannonEntropy(dataSet):
    numEntries = len(dataSet)
    labelCounts = {}
    for vec in dataSet:
        label = vec[-1]
        labelCounts[label] = labelCounts.get(label, 0) + 1
    
    shannonEntropy = 0.0
    for label in labelCounts:
        probability = float(labelCounts[label]) / numEntries
        shannonEntropy -= probability * log(probability, 2)
    return shannonEntropy

# 根据特征划分数据集
def splitDataSet(dataSet, fAxis, value):
    retDat = []
    for vec in dataSet:
        if vec[fAxis] == value:
            tmp = vec.copy() # 等价于 tmp = vec[:fAxis].extend(vec[fAxis+1:])
            del tmp[fAxis]
            retDat.append(tmp)
    return retDat

# 选择最优的数据集划分特征
def chooseBestFeature(dataSet):
    numDs = float(len(dataSet))
    numFeatures = len(dataSet[0]) - 1
    baseEntropy = calcShannonEntropy(dataSet)
    bestInfoGain = 0.0; bestFeature = -1
    for i in range(numFeatures):
        values = [example[i] for example in dataSet]
        valuesSet = set(values)
        newEntropy = 0.0
        for v in valuesSet:
            subDataV = splitDataSet(dataSet, i, v)
            prop = float(len(subDataV)) / numDs
            newEntropy += prop * calcShannonEntropy(subDataV)
            
        infoGain = baseEntropy - newEntropy
        print ('feature ', i, 'infoGain: ', infoGain)
        if infoGain >= bestInfoGain:
            bestFeature = i
            bestInfoGain = infoGain
    return bestFeature

# 对于未能完全划分的叶子节点根据投票来获取分类
def cleafMajorityCount(leafList):
    labelCount = {}
    for vote in leafList:
        labelCount[vote] = labelCount.get(vote, 0) + 1
    sortedCount = sorted(labelCount.items(), key = operator.itemgetter(1), reverse = True)
    return sortedCount[0][0]

# 递归创建决策树
def createDecisionTree(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 cleafMajorityCount(classList)
    
    bestFeature = chooseBestFeature(dataSet)
    bestFeatureLabel = labels[bestFeature]
    myTree = {bestFeatureLabel: {}}
    del labels[bestFeature]
    values = [example[bestFeature] for example in dataSet]
    valuesSet = set(values)
    for value in valuesSet:
        subLabels = labels[:]
        myTree[bestFeatureLabel][value] = createDecisionTree(splitDataSet(dataSet, bestFeature, value), subLabels)
    return myTree
    
            
# 创建数据集
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

import copy

# 熵函数测试
myDat, labels = createDataSet()
print (myDat)
print ('Initial Shannon Entropy: ', calcShannonEntropy(myDat), '\n')
'''
熵越高，则混合的数据也就越多，我们可以在数据集增加新的分类，观察熵的变化，这里增加一个新的 'maybe' 分类。
'''
testDat = copy.deepcopy(myDat)
testDat[0][-1] = 'maybe'
print (testDat)
print ('Add another label', calcShannonEntropy(testDat))

# 测试根据特征划分数据集函数 
print ('Split by feature 0(no surfacing)')
print(splitDataSet(myDat, 0, 0), '\t', splitDataSet(myDat, 0, 1))
print ('Split by feature 1(flippers)')
print(splitDataSet(myDat, 1, 0), '\t', splitDataSet(myDat, 1, 1))

# 测试选择最优的划分特征
print (myDat)
print ('Initial Choose:', chooseBestFeature(myDat), '\n')
print (splitDataSet(myDat, 0, 0))
print ('Another Choose:',chooseBestFeature(splitDataSet(myDat, 0, 0)))

# test cleafMajorityCount
cleafMajorityCount(['y','y','n'])

# 获取最终的决策树
myDat, labels = [[1, 'yes'], [1, 'yes'], [0, 'no']], ['flippers']
print (myDat, labels)
print (createDecisionTree(myDat, labels))
print ()
myDat, labels = [[1, 'yes'], [1, 'yes']], ['flippers']
print (myDat, labels)
print (createDecisionTree(myDat, labels))
print ()
myDat, labels = createDataSet()
createDecisionTree(myDat, labels)

'''
绘制决策树
'''
import matplotlib.pyplot as plt

# 获取树的深度
def getTreeDepth(tree):
    if tree == None or (not isinstance(tree, dict)) or len(tree.keys()) == 0:
        return 0
    subTreeKey = list(tree)[0]
    subTree = tree[subTreeKey]
    subTreeKeys = subTree.keys()
    maxSubTreeDepth = 0
    for subKey in subTreeKeys:
        depth = getTreeDepth(subTree[subKey])
        if depth >= maxSubTreeDepth:
            maxSubTreeDepth = depth
    return 1 + maxSubTreeDepth

# 获取树的最大宽度，即叶子节点个数。
def getTreeWidth(tree):
    if tree == None:
        return 0
    if isinstance(tree, str):
        return 1
    keys = list(tree.keys())
    subTree = tree[keys[0]]
    subTreeKeys = list(subTree)
    numLeaves = 0
    for subKey in subTreeKeys:
        numLeaves += getTreeWidth(subTree[subKey])
    return numLeaves

# 判断是否是叶子节点
def isLeaf(node):
    return isinstance(node, dict) and len(node.keys()) == 1 and isinstance(node[list(node)[0]], str)

# 获取树的根节点位置
def getRootPos(tree, xStartPos = 0):
    width = getTreeWidth(tree)
    height = getTreeDepth(tree)
    return (width - 1) * 4 + xStartPos, height * 2

# 绘制树的节点
def plotTreeNodes(tree, xStartPos = 0):
    if isLeaf(tree):
        print(xStartPos + 2, 2)
        return
    xPos, yPos = getRootPos(tree)
    keys = list(tree.keys())
    rootTree = tree[keys[0]]
    subKeys = list(rootTree)
    xStartPos = 0
    for subKey in subKeys:
        print ('subTree:', rootTree, subKey)
        subTree = rootTree[subKey]
        plotTreeNodes(subTree, xStartPos)
        xStartPos, y = getRootPos(subTree)
    print (xPos, yPos)

tree1 = {'no surfacing': {0: 'no', 1: {'flippers': {0: {'flippers': {0: 'no', 1: 'yes'}}, 1: 'yes'}}}}
tree2 = {'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}}
tree3 = {'no surfacing': {0: 'no', 1: 'no', 2: 'yes'}}
print(getTreeDepth(tree1), getTreeDepth(tree2), getTreeDepth(tree3))
print(getTreeWidth(tree1), getTreeWidth(tree2), getTreeWidth(tree3))
print(isLeaf({0: 'yes'}), isLeaf(tree3))
print(getRootPos(tree1), getRootPos(tree2), getRootPos(tree3))
print ('Plot tree')
plotTreeNodes(tree1)