基于python的sklearn库的决策树算法基本实现

基于python的sklearn库的决策树算法基本实现

关于不同年龄等几个特征的人进行是否购买电脑的预测

先贴代码

# -*- coding:utf-8 -*-

"""
fit()可以说是调用的通用方法
fit(X)，表示用数据X来训练某种模型。 函数返回值一般为调用fit方法的对象本身
fit(X,y=None)为无监督学习算法，fit(X,Y)为监督学习算法
"""

import csv
from sklearn.feature_extraction import DictVectorizer  # 字典元素的list
from sklearn import preprocessing
from sklearn import tree

File = open(r"E:\pyTree.csv", "rb")  # 文件读写 见下link
reader = csv.reader(File)  # 读取csv文件 返回<type '_csv.reader'>
head = reader.next()  # next 类似与指针后移一行 所以这一般是用来读标题行(head)

resultList = []  # 预测的结果列表
featureList = []  # 特征值 = 标题:数值 (字典的集合)

for line in reader:
    resultList.append(line[-1])  # 一行中的具体位置具体修改
    dic = {}  # 临时字典
    for row in range(1, len(line) - 1):
        dic[head[row]] = line[row]  # 取得各个特征
    featureList.append(dic)

vec = DictVectorizer()
binx = vec.fit_transform(featureList).toarray()
# fit_transform函数 先调用fit(获取特征信息和目标值信息)  后调用transform(对特征进行转换) 来训练模型

lb = preprocessing.LabelBinarizer()  # 二进制化
biny = lb.fit_transform(resultList)  # 预测的结果列表 二值化


clf = tree.DecisionTreeClassifier(criterion='entropy')  # 不带参数 默认gini
clf = clf.fit(binx, biny)  # data, label

with open("out.dot", "w") as f:  # 打印决策树文件 不做分析
    f = tree.export_graphviz(clf, feature_names=vec.get_feature_names(), out_file=f)

# 简单测试
testLine = binx[0, :]
testLine[0] = 1
testLine[2] = 0

ans = clf.predict([testLine, binx[0, :]])  # 预测格式应该与读入特征值数据的格式相同 本例是二维list
print ans  # 10代表是否购买电脑