Python的机器学习包中的DecisionTreeClassifier可以进行决策树分类。
1.输入数据集。
import pandas as pddf=pd.DataFrame({'name':['Lily','Lucy','Jim','Tom','Anna','Jack','Sam'],'weight':[42,38,78,67,52,80,92],'height':[162,158,169,170,166,175,178],'is_fat':[0,0,1,0,1,0,1]})
2.导入决策树工具包。
from sklearn.tree import DecisionTreeClassifier
3.准备训练集,is_fat为目标变量,'weight'和'height'为自变量。
X=df.loc[:,['weight','height']]
y=df['is_fat']
4.建立模型,并进行模型训练。
clf=DecisionTreeClassifier()
clf.fit(X,y)
5.利用模型进行预测。
y_pred=clf.predict(X)
print(y_pred)
6.根据预测结果绘制散点图。
import matplotlib.pyplot as plt
plt.figure()
df['is_fat_pred']=y_pred
df_0=df[df['is_fat_pred']==0]
df_1=df[df['is_fat_pred']==1]
plt.scatter(df_0['weight'],df_0['height'],c='y',s=50,label='normal')
plt.scatter(df_1['weight'],df_1['height'],c='lightblue',s=100,label='fat')
for k in range(len(X)):
plt.text(X['weight'][k],X['height'][k],df['name'][k]) #设置散点标签
plt.legend()
plt.show()