Trees/trees.py at master · SXKai/Trees · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
# -*- coding: utf-8 -*-
"""
Created on Sun Oct 29 15:54:15 2017

@author: Q
"""
import numpy as np
import math
import operator
def createDataSet():
    dataSet = [[1,1,'yes'],
               [1,1,'yes'],
               [1,0,'no'],
               [0,1,'no'],
               [0,1,'no'],]
    label = ['no surfacing','flippers']
    return dataSet,label

def calcShannonEnt(dataSet): #计算信息熵
    numEntries = len(dataSet)
    labelCounts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        if currentLabel not in labelCounts:
            labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1
    shannonEnt = 0
    for key in labelCounts:
        shannonEnt = shannonEnt - (labelCounts[key]/numEntries)*math.log2(labelCounts[key]/numEntries)
    return shannonEnt


def splitDataSet(dataSet,axis,value):#切割数据集
    retDataset = []
    for featVec in dataSet:
        if featVec[axis] == value:
            newVec = featVec[:axis]
            newVec.extend(featVec[axis+1:])
            retDataset.append(newVec)
    return retDataset

def chooseBestFeatureToSplit(dataSet):#选择最优特征
    numFeatures = len(dataSet[0]) - 1
    bestInfoGain = 0
    bestFeature = -1
    baseEntropy = calcShannonEnt(dataSet)
    for i in range(numFeatures):
        allValue = [example[i] for example in dataSet]
        allValue = set(allValue)
        newEntropy = 0
        for value in allValue:
            splitset = splitDataSet(dataSet,i,value)
            newEntropy = newEntropy + len(splitset)/len(dataSet)*calcShannonEnt(splitset)
        infoGain = baseEntropy - newEntropy
        if infoGain > bestInfoGain:
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature

def majorityCnt(classList):#投票确定叶子节点的类
    classCount = {}
    for value in classList:
        if value not in classCount: classCount[value] = 0
        classCount[value] += 1
    classCount = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True)
    return classCount[0][0]

def createTree(dataSet,labels):   #生成决策树
    classList = [example[-1] for example in dataSet]
    labelsCopy = labels[:]
    if classList.count(classList[0]) == len(classList):
        return classList[0]
    if len(dataSet[0]) == 1:
        return majorityCnt(classList)
    bestFeature = chooseBestFeatureToSplit(dataSet)
    bestLabel = labelsCopy[bestFeature]
    myTree = {bestLabel:{}}
    featureValues = [example[bestFeature] for example in dataSet]
    featureValues = set(featureValues)
    del(labelsCopy[bestFeature])
    for value in featureValues:
        subLabels = labelsCopy[:]
        myTree[bestLabel][value] = createTree(splitDataSet(dataSet,bestFeature,value),subLabels)
    return myTree


def classify(inputTree,featLabels,testVec):
    currentFeat = list(inputTree.keys())[0]
    secondTree = inputTree[currentFeat]
    try:
        featureIndex = featLabels.index(currentFeat)
    except ValueError as err:
        print('yes')
    try:
        for value in secondTree.keys():
            if value == testVec[featureIndex]:
                if type(secondTree[value]).__name__ == 'dict':
                    classLabel = classify(secondTree[value],featLabels,testVec)
                else:
                    classLabel = secondTree[value]
        return classLabel
    except AttributeError:
        print(secondTree)

if __name__ == "__main__":
    dataset,label = createDataSet()
    myTree = createTree(dataset,label)
    a = [1,1]
    print(classify(myTree,label,a))