决策树算法简单应用

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

from sklearn import tree

# visualize code
from sklearn.externals.six import StringIO
import pydotplus

# 决策树算法
# 初步的两个特性的判断，[重量，表皮光滑度](对于水果，可以是：1=光滑，0=粗糙)
# 结论标签，1=苹果，0=橘子
features = [[140,1],[130,1],[150,0],[170,0],[150,1],[130,0]]
features_names = ['重量','表皮光滑度']
labels   = [0, 0, 1, 1, 0, 1]
label_name = ['橘子','苹果']

#调用决策树算法的核心语句
dt = tree.DecisionTreeClassifier()
dt= dt.fit(features, labels)

#测试数据，预测[200,1]
print(dt.predict([[200,1]]))
# 可以根据测试数据，得到预测的结果

# 此处添加三个"截止
# 移动标记来进行代码运行测试
dot_data = StringIO()
tree.export_graphviz(dt,
                        out_file = dot_data,
                        feature_names = features_names,
                        class_names   = label_name,
                        filled   = True, rounded = True,
                        impurity = False)
graph = pydotplus.graph_from_dot_data(dot_data.getvalue().replace('helvetica','"Microsoft YaHei"'))
print(dot_data)

#graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
#下面解决Windows下乱码问题
graph = pydotplus.graph_from_dot_data(dot_data.getvalue().replace('helvetica','"Microsoft YaHei"'))
print(graph)
graph.write_pdf("basic.pdf")

从执行结果可以知道预测[200,1]为苹果

runfile('D:/ai/py/tensorflow-basic.py', wdir='D:/ai/py')
[0]
<_io.StringIO object at 0x000001D93B47CA68>
<pydotplus.graphviz.Dot object at 0x000001D93B905F98>

修改数据

features = [[140,1],[130,1],[150,0],[170,0],[150,1],[130,0]]
features_names = ['重量','表皮光滑度']
labels   = [0, 0, 1, 0, 0, 1]

则[200,0]的预测结果是[0]