Skip to main content
Back to Blog
AI/MLData AnalysisProgramming Languages
13 August 20264 min readUpdated 13 August 2026

Implementing Decision Tree Regression in Python with Sklearn

A decision tree regressor is a machine learning model used for predicting continuous outcomes, such as prices or scores, through a structure that mimics a tree. By segmenting da...

Implementing Decision Tree Regression in Python with Sklearn

A decision tree regressor is a machine learning model used for predicting continuous outcomes, such as prices or scores, through a structure that mimics a tree. By segmenting data into smaller sets based on straightforward rules derived from the input features, it can effectively minimize prediction errors. The model provides a prediction at each leaf node, typically the average value of that segment.

For instance, when predicting house prices using factors like size, location, and age, the tree might initially split by location, followed by size, and then age. Here's how to implement it:

Step 1: Importing the Required Libraries

To begin, import the necessary libraries:

  • NumPy for handling numerical computations and arrays.
  • Matplotlib for creating graphs and visualizations.
  • Various modules from scikit-learn for tasks like modeling, data splitting, tree visualization, and performance evaluation.
import numpy as np
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeRegressor, export_text
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

Step 2: Creating a Sample Dataset

A synthetic dataset is created using NumPy, where feature values X are generated randomly and sorted between 0 and 5. The target y is a noisy sine function of X. A scatter plot shows the relationship between the feature and target values.

np.random.seed(42)
X = np.sort(5 * np.random.rand(100, 1), axis=0)
y = np.sin(X).ravel() + np.random.normal(0, 0.1, X.shape[0])

plt.scatter(X, y, color='red', label='Data')
plt.title("Synthetic Dataset")
plt.xlabel("Feature")
plt.ylabel("Target")
plt.legend()
plt.show()

Step 3: Splitting the Dataset

The dataset is divided into training and testing subsets using the train_test_split function, with 70% allocated for training and 30% for testing. Setting random_state=42 ensures reproducibility.

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

Step 4: Initializing the Decision Tree Regressor

A decision tree regressor is initialized with a maximum depth of 4 to control the complexity of the model.

regressor = DecisionTreeRegressor(max_depth=4, random_state=42)

Step 5: Fitting the Decision Tree Regressor Model

The model is trained using the .fit() method on the training data (X_train and y_train), allowing it to learn the relationships between variables.

regressor.fit(X_train, y_train)

Step 6: Predicting a New Value

The trained model is used to predict new values with the predict() function. The mean squared error (MSE) is calculated to evaluate the model's accuracy on unseen test data.

y_pred = regressor.predict(X_test)

mse = mean_squared_error(y_test, y_pred)
print(f"Mean Squared Error: {mse:.4f}")

Step 7: Visualizing the Result

The model's predictions are visualized to assess how well the decision tree fits the data. This visualization shows the step-like segments of predictions based on the tree’s splits.

X_grid = np.arange(min(X), max(X), 0.01)[:, np.newaxis]
y_grid_pred = regressor.predict(X_grid)

plt.figure(figsize=(10, 6))
plt.scatter(X, y, color='red', label='Data')
plt.plot(X_grid, y_grid_pred, color='blue', label='Model Prediction')
plt.title("Decision Tree Regression")
plt.xlabel("Feature")
plt.ylabel("Target")
plt.legend()
plt.show()

Step 8: Exporting and Displaying the Tree Structure

The plot_tree function is used to visualize the decision tree structure, illustrating how the model splits the feature space and partitions the data for predictions.

from sklearn.tree import plot_tree

plt.figure(figsize=(20, 10))
plot_tree(
    regressor,
    feature_names=["Feature"],
    filled=True,
    rounded=True,
    fontsize=10
)
plt.title("Decision Tree Structure")
plt.show()

Decision tree regression is a powerful tool for predicting continuous values, capable of capturing non-linear patterns in data. Its interpretability makes it easy to understand the decision-making process of the model.