《统计学习方法》第五章——回归树 python实现

import numpy as np

def divide(data):
    divide_x = []
    divide_y = []
    for i in range(len(data) - 1):
        divide_x.append(np.mean([data[i][0], data[i + 1][0]]))
    for i in range(len(data) - 1):
        y1_arr = []
        y2_arr = []
        for j in range(len(data)):
            if j + data[0][0] < divide_x[i]:
                y1_arr.append(data[j][1])
            else:
                y2_arr.append(data[j][1])
        y_temp = [y1_arr, y2_arr]
        divide_y.append(y_temp)
    return divide_x,divide_y

def min_mse(divide_x,divide_y):
    mse_arr = []
    for i in range(len(divide_x)):
        ave_y1 = np.mean(divide_y[i][0])
        ave_y2 = np.mean(divide_y[i][1])
        mse = 0
        for j in range(len(divide_y[i][0])):
            mse += np.square(divide_y[i][0][j] - ave_y1)
        for k in range(len(divide_y[i][1])):
            mse += np.square(divide_y[i][1][k] - ave_y2)
        mse_arr.append(mse)

    min_mse = 10000
    index = 10000
    for i in range(len(mse_arr)):
        if min_mse > mse_arr[i]:
            min_mse = mse_arr[i]
            index = i
    return index

class BinaryTree:
    def __init__(self,prop):
        self.prop=prop
        self.leftchild=None
        self.rightchild=None
        self.value=None

data=[[1,4.50],[2,4.75],[3,4.91],[4,5.34],[5,5.80],[6,7.05],[7,7.90],[8,8.23],[9,8.70],[10,9.00]]

def create_tree(data,prop):
    if len(data)==1:
        return
    temp_tree_value = [prop, data]
    root = BinaryTree(prop)
    root.value = temp_tree_value
    divide_x, divide_y = divide(data)
    index = min_mse(divide_x, divide_y)
    print(divide_y[index])
    create_tree(data[:index + 1],"left")
    create_tree(data[index + 1:len(data)],"right")

create_tree(data,"root")