python使用sklearn实现决策树的方法示例

# -*- coding: utf-8 -*-"""Created on Sun Dec 25 11:25:40 2016@author: Administrator"""from sklearn.feature_extraction import DictVectorizerimport csvfrom sklearn import treefrom sklearn import preprocessingfrom sklearn.externals.six import StringIOimport pydotplusfrom IPython.display import Image# Read in the csv file and put features into list of dict and list of class labelallElectornicsData = open('AllElectronics.csv', 'r')reader = csv.reader(allElectornicsData)# headers = reader.next()  python2.7 supported  本质获取csv 文件的第一行数据#headers = reader.__next__()  python 3.5.2 headers = next(reader)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)print(labelList)# Vetorize featuresvec = DictVectorizer()dummyX = vec.fit_transform(featureList).toarray()print("dummyX: " + str(dummyX))print(vec.get_feature_names())print("labelList: " + str(labelList))# vectorize class labelslb = preprocessing.LabelBinarizer()dummyY = lb.fit_transform(labelList)print("dummyY: ", str(dummyY))# Using decision tree for classification    ===========【此处调用为算法核心】============#clf = tree.DecisionTreeClassifier(criterion='entropy')clf = tree.DecisionTreeClassifier(criterion='gini')clf = clf.fit(dummyX, dummyY)print("clf: ", str(clf))# Visualize model# dot -Tpdf iris.dot -o ouput.pdfwith open("allElectronicInformationGainOri.dot", 'w') as f:  f = tree.export_graphviz(clf, feature_names = vec.get_feature_names(), out_file = f)# predictoneRowX = dummyX[0, :]print("oneRowX: " + str(oneRowX))newRowX = oneRowXnewRowX[0] = 1newRowX[2] = 0print("newRowX: " + str(newRowX))predictedY = clf.predict(newRowX)print("predictedY: " + str(predictedY))