Python机器学习：决策树001什么是决策树

最新推荐文章于 2024-10-15 17:32:38 发布

范德彪陕西分彪

最新推荐文章于 2024-10-15 17:32:38 发布

阅读量216

点赞数

分类专栏： Python机器学习

本文链接：https://blog.csdn.net/weixin_46815330/article/details/111630916

版权

Python机器学习专栏收录该内容

73 篇文章 3 订阅

订阅专栏

在这里插入图片描述

import numpy as np
import matplotlib.pyplot as plt

from sklearn import datasets
iris = datasets.load_iris()
X = iris.data[:,2:]
y = iris.target

from sklearn.tree import DecisionTreeClassifier
dt_clf = DecisionTreeClassifier(max_depth = 2,criterion = 'entropy')#熵
dt_clf.fit(X,y)

def plot_decision_boundary(model, axis):

    x0, x1 = np.meshgrid(
        np.linspace(axis[0], axis[1], int((axis[1]-axis[0])*100)).reshape(-1, 1),
        np.linspace(axis[2], axis[3], int((axis[3]-axis[2])*100)).reshape(-1, 1),
    )
    X_new = np.c_[x0.ravel(), x1.ravel()]

    y_predict = model.predict(X_new)
    zz = y_predict.reshape(x0.shape)

    from matplotlib.colors import ListedColormap
    custom_cmap = ListedColormap(['#EF9A9A','#FFF59D','#90CAF9'])

    plt.contourf(x0, x1, zz, linewidth=5, cmap=custom_cmap)

plot_decision_boundary(dt_clf,axis=[0.5,7.5,0,3])
plt.scatter(X[y == 0,0],X[y == 0,1])
plt.scatter(X[y == 1,0],X[y == 1,1])
plt.scatter(X[y == 2,0],X[y == 2,1])