Decision tree regression

In [1]:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

In [2]:

dataset = pd.read_csv('Position_Salaries.csv')
dataset.head(2)

Out[2]:

	Position	Level	Salary
0	Business Analyst	1	45000
1	Junior Consultant	2	50000

In [3]:

x = dataset.iloc[:,1:2].values
y = dataset.iloc[:,2].values

In [4]:

from sklearn.tree import DecisionTreeRegressor
reg = DecisionTreeRegressor()
reg.fit(x,y)
plt.scatter(x,y)
plt.plot(x,reg.predict(x))
reg.predict(6.5)

Out[4]:

array([150000.])

In [5]:

x_grid = np.arange(min(x),max(x),0.01)
x_grid = x_grid.reshape(len(x_grid),1)
plt.scatter(x,y)
plt.plot(x_grid,reg.predict(x_grid))

Out[5]:

[<matplotlib.lines.Line2D at 0x1c7c1315c18>]

In [ ]: