Python实现基于决策树的西瓜分类

from sklearn.feature_extraction import DictVectorizer
from sklearn import tree
from sklearn import preprocessing
import csv
import graphviz

# 读取西瓜数据集
Dtree = open('西瓜数据集3.0.csv', 'r')
reader = csv.reader(Dtree)

# 获取特征名称
headers = reader.__next__()
print(headers)

# 存储特征和标签数据
featureList = []
labelList = []

# 遍历数据集，提取特征和标签
for row in reader:
    # 提取标签
    labelList.append(row[-1])
    # 提取特征
    rowDict = {}
    for i in range(1, len(row)-3):
        rowDict[headers[i]] = row[i]
    featureList.append(rowDict)
print(featureList)

# 特征向量化
vec = DictVectorizer()
x_data = vec.fit_transform(featureList).toarray()
print('x_data: ' + str(x_data))

# 标签二进制编码
lb = preprocessing.LabelBinarizer()
y_data = lb.fit_transform(labelList)
print('y_data: ' + str(y_data))

# 创建决策树模型
model = tree.DecisionTreeClassifier(criterion='entropy')
# 训练模型
model.fit(x_data, y_data)

# 测试模型
x_test = x_data[0]
predict = model.predict(x_test.reshape(1,-1))
print('predict: ' + str(predict))

# 决策树可视化
dot_data = tree.export_graphviz(model,
                                out_file=None,
                                feature_names=vec.get_feature_names(),
                                class_names=lb.classes_,
                                filled=True,
                                rounded=True,
                                special_characters=True)
graph = graphviz.Source(dot_data)
graph.render('Tree')

代码解释：

导入必要的库: 包括DictVectorizer用于特征向量化，tree用于决策树模型，preprocessing用于数据预处理，csv用于读取CSV文件，graphviz用于决策树可视化。
读取数据: 使用csv.reader读取CSV文件中的数据，并将特征名称存储在headers列表中。
数据预处理: 将特征存储在featureList列表中，将标签存储在labelList列表中。
特征工程: 使用DictVectorizer将特征字典转换为特征向量，使用LabelBinarizer将标签转换为二进制编码。
模型训练: 创建决策树模型，并使用训练数据进行训练。
模型预测: 使用测试数据进行预测，并将预测结果打印出来。
模型可视化: 使用graphviz库将决策树可视化，并保存为图片。

代码改进:

使用更具描述性的变量名。
添加注释以解释代码的功能。
将代码封装成函数，以便于复用。
使用交叉验证来评估模型的性能。