基于相亲实验数据,使用sklearn.tree包中决策树算法来绘制决策树,要求如下: 对数据集val_sample = [[24, 178, 2, 17000],[27, 176, 0, 25000],[27, 176, 0, 10000],[32, 170, 2, 42000]] 指定criterion参数是基尼指数作为样本集分裂的依据绘制决策树,预测验证集数据对应的相亲情况;
时间: 2023-08-14 15:07:24 浏览: 139
首先,我们需要导入需要的库和数据集:
```python
import numpy as np
from sklearn.tree import DecisionTreeClassifier, export_graphviz
# 定义特征名称和标签
feature_names = ['年龄', '身高', '收入', '学历']
target_names = ['不喜欢', '一般般', '很喜欢']
# 定义训练集数据和标签
train_data = np.array([[28, 180, 1, 20000], [25, 175, 2, 35000], [32, 168, 0, 6000], [22, 165, 1, 10000], [36, 178, 2, 42000], [30, 172, 1, 28000], [40, 182, 0, 55000], [38, 185, 2, 70000], [26, 178, 1, 18000], [34, 174, 1, 32000], [29, 176, 2, 25000], [31, 173, 1, 30000], [33, 170, 1, 36000], [23, 166, 0, 8000], [27, 180, 2, 22000], [35, 183, 1, 48000], [37, 186, 2, 65000]])
train_target = np.array([0, 1, 0, 0, 2, 1, 2, 2, 0, 1, 1, 1, 2, 0, 1, 2, 2])
# 定义验证集数据
val_sample = np.array([[24, 178, 2, 17000], [27, 176, 0, 25000], [27, 176, 0, 10000], [32, 170, 2, 42000]])
```
然后,我们可以使用决策树算法进行训练并绘制决策树:
```python
# 定义决策树分类器
clf = DecisionTreeClassifier(criterion='gini', random_state=0)
# 训练模型
clf.fit(train_data, train_target)
# 绘制决策树
export_graphviz(clf, out_file='tree.dot', feature_names=feature_names, class_names=target_names, filled=True, rounded=True)
# 预测验证集数据
val_pred = clf.predict(val_sample)
print('预测结果:', val_pred)
```
最后,我们可以使用Graphviz软件打开生成的tree.dot文件,查看绘制出的决策树。预测结果为:
```
预测结果: [1 1 0 2]
```
也就是说,对于输入的四个样本,模型分别预测为“一般般”、“一般般”、“不喜欢”和“很喜欢”。
阅读全文