决策树

from sklearn.feature_extraction import DictVectorizer
import csv
from sklearn import preprocessing
from sklearn import tree
from sklearn.externals.six import StringIO
import graphviz

#read in the csv file and put features in a list of dict
data = open('DecisionTreeTrainData.csv')
reader = csv.reader(data)
headers = next(reader) #python2中用header = reader.next()

print(headers)

featureList = []
labelList = []

for row in reader:
    labelList.append(row[len(row) - 1])
    rowDict =  {}
    for i in range(1, len(row) - 1):
        rowDict[headers[i]] = row[i]
    
    featureList.append(rowDict)

print(featureList)

#Vectorize features
vec = DictVectorizer()
dummyX = vec.fit_transform(featureList).toarray() #将每一个样本四种属性特征的十种可能情况，用0或1的十维向量表示出来

print('dummyX:' + str(dummyX)) #输出十四个样本，每个样本由十维向量表示出来
print(vec.get_feature_names())

print('labelList:' + str(labelList))

#Vectorize class labels
lb = preprocessing.LabelBinarizer()
dummyY = lb.fit_transform(labelList)
print('dummyY:' + str(dummyY))

#using decision tree for classification
clf = tree.DecisionTreeClassifier(criterion = 'entropy') #选用ID3方法，利用信息熵作为度量标准
clf = clf.fit(dummyX, dummyY)
print('clf:' + str(clf))

#Visulize model
with open('DecisionTreeTrainData.dot', 'w') as f:
    f = tree.export_graphviz(clf, feature_names = vec.get_feature_names(), out_file = f)

#将dot文件通过graphviz输出成pdf文件
with open('DecisionTreeTrainData.dot') as gf:
    dot_graph = gf.read()
    dot = graphviz.Source(dot_graph) #调用graphviz的source方法读取文件，然后调用view方法显示视图
    dot.view()

##上述已经将csv文件中数据装换成决策树树图，下面将第一行用户数据更改后作为test数据进行预测    
oneRowX = dummyX[0,:]
print('oneRowX:' + str(oneRowX))

newRowX = oneRowX
newRowX[0] = 1
newRowX[2] = 0

print('newRowX:' + str(newRowX))


predictedY = clf.predict([newRowX])
print('predictedY:' + str(predictedY))