python cart树_CART回归树python实现

import numpy as np

def reg_leaf(dataSet):

"""生成叶子节点，返回这个叶子上样本标记的平均值"""

return np.mean(dataSet[:, -1])

def reg_err(dataSet):

"""计算平方误差"""

return np.var(dataSet[:, -1]) * shape(dataSet)[0]

def bin_split_data(dataSet, feat_index, split_value):

"""二分数据集"""

arr0 = []

arr1 = []

for i in range(shape(dataSet)[0]):

if dataSet[i, feat_index] <= split_value:

arr0.append(dataSet[i])

else:

arr1.append(dataSet[i])

return np.mat(arr0), np.mat(arr1)

def choose_best_split(dataSet, leaf_type=reg_leaf, err_type = reg_err, ops=(1,4)):

"""选择最优分裂节点和分裂值"""

tolS = ops[0] #误差减少阈值，达到tolS才允许分裂

tolN = ops[1] #最少分割样本，达到tolN才允许分裂

if len(set(dataSet[:,-1])) == 1 or len(dataSet) <= tolN: #样本的值相等，没必要分割了

return None, leaf_type(dataSet)

m, n = shape(dataSet)

ori_err = reg_err(dataSet)

lowest_err = np.inf

best_index = -1

best_value = -1

for index in range(n - 1):

for value in set(dataSet[:, index]):

mat0, mat1 = bin_split_data(dataSet, index, value)

if shape(mat0)[0] < tolN or shape(mat1) < tolN: #子节点样本数过少

continue

new_err = reg_err(mat0, mat1) #子树的平方误差和

if new_err < lowest_err:

lowest_err = new_err

best_index = index

best_value = value

if best_index == -1 or (ori_err - lowest_err) < tolS: #误差减少太小

return None, leaf_type(dataSet)

return best_index, best_value

def create_tree(dataSet, feat_names):

"""递归创建一棵决策树"""

y_train = [row[-1] for row in dataSet]

if len(set(y_train)) == 1: #类别完全相同，停止继续划分，返回类别

return y_train[0]

if len(dataSet[0]) == 1: #没有特征可以划分了，直接返回最多的特征

return get_majority_cnt_label(y_train)

best_feat = choose_best_feat_ent_ratio(dataSet) #找到最优分割特征

best_feat_name = feat_names[best_feat]

myTree = {best_feat_name:{}} #开始构建二叉树

del feat_names[best_feat]

feat_data_dict = splitDataSet(col_id=best_feat, dataSet=dataSet)

for feat_value, data in feat_data_dict.items():

sub_feat_names = feat_names[:] #拷贝赋值，防止被修改

myTree[best_feat_name][feat_value] = create_tree(data, sub_feat_names) #保证传进去的不是空的数据集

return myTree