python 实现CART算法决策树

有用请点赞，没用请差评。

欢迎分享本文，转载请保留出处。

 

本次代码是基于上一节决策树ID3\C45修改过来的，建议两篇博客一起看。具体算法原理等有时间了再写。

 

 

# -*- coding:utf-8 -*-
# Decision tree by cart决策树，cart算法，算法参考李航《统计学习方法》P71
#author:Tomator


import numpy as np
import math
from sklearn.model_selection import train_test_split


# 测试数据集，《统计学习方法》P59，贷款申请样本数据集
def createDataSet():
    dataSet = [[0, 0, 0, 0, 'no'],  # 数据集
               [0, 0, 0, 1, 'no'],
               [0, 1, 0, 1, 'yes'],
               [0, 1, 1, 0, 'yes'],
               [0, 0, 0, 0, 'no'],
               [1, 0, 0, 0, 'no'],
               [1, 0, 0, 1, 'no'],
               [1, 1, 1, 1, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [2, 0, 1, 2, 'yes'],
               [2, 0, 1, 1, 'yes'],
               [2, 1, 0, 1, 'yes'],
               [2, 1, 0, 2, 'yes'],
               [2, 0, 0, 0, 'no']]
    labels = ['年龄', '有工作', '有自己的房子', '信贷情况']  # 分类属性
    return dataSet, labels  # 返回数据集和分类属性



# 计算基尼指数
def cal_gini(data_vector):
    nums_data = len(data_vector)  # 数据集样本数
    counts_by_labels = {}  # 用来保存每个label下的样本数
    gini = 0      #基尼指数
    p_sum=0       #每个类别的样本数

    for vector in data_vector:
        if vector[-1] not in counts_by_labels:  # vector[-1]为label值
            counts_by_labels[vector[-1]] = 0
        counts_by_labels[vector[-1]] +