py 决策树

from cgi import print_form
from math import log
from msilib import datasizemask
import os
from matplotlib import pyplot as plt
import numpy as np

import operator
import csv
from sklearn import *
from sklearn.tree import DecisionTreeClassifier
def loaddata ():
    dataSet = [[0, 0,0,0,0,0, 'yes'],
               [1, 0,1,0,0,0,'yes'],
               [1, 0,0,0,0,0,'yes'],
               [0, 0,1,0,0,0,'yes'],
               [2, 0,0,0,0,0,'yes'],
               [0, 1,0,0,1,1,'yes'],
               [1, 1,0,1,1,1,'yes'],
               [1, 1,0,0,1,0, 'yes'],
               [1, 1,1,1,1,0,'no'],
               [0, 2,2,0,2,1,'no'],
               [2, 2,2,2,2,0,'no'],
               [2, 0,0,2,2,1,'no'],
               [0, 1,0,1,0,0, 'no'],
               [2, 1,1,1,0,0,'no'],
               [1, 1,0,0,1,1,'no'],
               [2, 0,0,2,2,0,'no'],
               [0, 0,1,1,1,0,'no']]
    feature_name = ['a1','a2','a3','a4','a5','a6']
    return dataSet, feature_name

def loaddata_new():
    # 定义文件路径
    csv_path = 'watermelon2.csv'
    with open(csv_path,'r',encoding='utf-8-sig')as fp:
        dataSet = [i for i in csv.reader(fp)]  # csv.reader 读取到的数据是list类型
    feature_name = ['a1','a2','a3','a4','a5','a6']
    return dataSet, feature_name

def entropy(dataSet):
    #数据集条数
    m = len(dataSet)
#保存所有的类别及属于该类别的样本数
    labelCounts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        if currentLabel not in labelCounts.keys(): 
            labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1
    #保存熵值
    e = 0.0 
    #补充计算信息熵的代码
    for key in labelCounts:
        prob = float(labelCounts[key]) / m
        e -= prob * log(prob,2)
    return e

def splitDataSet(dataSet, axis, value):
    #创建返回的数据集列表
    retDataSet=[]
    #遍历数据集
    for featVec in dataSet:
        if featVec[axis]==value:
            #去掉axis特征
            reduceFeatVec=featVec[:axis]
            #将符合条件的添加到返回的数据集
            reduceFeatVec.extend(featVec[axis+1:])
            retDataSet.append(reduceFeatVec)
    #返回划分后的数据集

    return retDataSet

def chooseBestFeature(dataSet):
    n = len(dataSet[0]) - 1
    #计数整个数据集的熵
    baseEntropy = entropy(dataSet)
    bestInfoGain = 0.0; bestFeature = -1
    #遍历每个特征
    for i in range(n):  
        #获取当前特征i的所有可能取值
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList) 
        newEntropy = 0.0
        #遍历特征i的每一个可能的取值
        for value in uniqueVals:
            #按特征i的value值进行数据集的划分
            subDataSet = splitDataSet(dataSet, i, value)
            #补充计算条件熵的代码
            prob = len(subDataSet) / float(len(dataSet))
            newEntropy += prob * entropy((subDataSet))
        #计算信息增益
        infoGain = baseEntropy - newEntropy  
        #保存当前最大的信息增益及对应的特征
        if (infoGain > bestInfoGain):
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature

def classVote(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]

def trainTree(dataSet,feature_name):
    classList = [example[-1] for example in dataSet]
#所有类别都一致
    if classList.count(classList[0]) == len(classList): 
        return classList[0] 
#数据集中没有特征
    if len(dataSet[0]) == 1: 
        return classVote(classList)
#选择最优划分特征
    bestFeat = chooseBestFeature(dataSet)
    bestFeatName = feature_name[bestFeat]
    myTree = {bestFeatName:{}}
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues)
#遍历uniqueVals中的每个值，生成相应的分支
    for value in uniqueVals:
        sub_feature_name = feature_name[:]
        # 生成在dataSet中bestFeat取值为value的子集；
        sub_dataset = splitDataSet(dataSet,bestFeat,value)      #补充代码
        # 根据得到的子集，生成决策树
        myTree[bestFeatName][value] = trainTree(sub_dataset,feature_name)      #补充代码
    return myTree

def tree3(clf):
    fig = plt.figure(figsize=(35, 10))
    tree.plot_tree(clf, fontsize=8)
    fig.savefig(os.path.join(".\\", "tree3.png"))

myDat,feature_name = loaddata()

t = np.array(myDat)
myTree = trainTree(myDat,feature_name)
tree3(myTree)
print(myTree)

# print(t)
# x = t[:,:-1]
# y = t[:,-1]
# tr = DecisionTreeClassifier(criterion="entropy")
# tr.fit(x,y)
# tree3(tr)
# tree.plot_tree(tr)

# print(x)
# print(y) 
# myDat,feature_name = loaddata()
# myTree = trainTree(myDat,feature_name)
# print(myTree)