参考链接: Python | 使用sklearn进行决策树回归
使用决策树回归模型 DecisionTreeRegressor 来完成回归预测在波士顿房价数据集进行操作
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler,MinMaxScaler
from sklearn.datasets import load_boston
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
data = load_boston()
x = data.data
y = data.target
df = pd.DataFrame(x,columns=data.feature_names)
df['price'] = y
corr = df.corr()['price']
corr = corr.abs()
corr1 = corr.copy().drop('price')
features = []
n = 3 # 使用相关性最大的前三个特征
for i in range(n):
index = corr1.argmax()
features.append(corr1.keys()[index])
corr1 = corr1.drop(corr1.keys()[index])
print(features)
new_x = df[features]
scale = MinMaxScaler()
new_x = scale.fit_transform(new_x)
x_train,x_test,y_train,y_test = train_test_split(new_x,y,test_size=0.33,random_state=17)
model = DecisionTreeRegressor()
model.fit(x_train,y_train)
score = model.score(x_test,y_test)