import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
df = pd.read_csv('Position_Salaries.csv')
X = df.iloc[:, 1:2] # Using 1:2 as indices will give us np array of dim (10, 1)
y = df.iloc[:, 2]
df.head()
Position | Level | Salary | |
---|---|---|---|
0 | Business Analyst | 1 | 45000 |
1 | Junior Consultant | 2 | 50000 |
2 | Senior Consultant | 3 | 60000 |
3 | Manager | 4 | 80000 |
4 | Country Manager | 5 | 110000 |
from sklearn.tree import DecisionTreeRegressor
regressor = DecisionTreeRegressor(random_state=0).fit(X, y)
plt.scatter(X, y)
X_grid = np.arange(min(X.values), max(X.values), 0.1).reshape(-1, 1)
plt.plot(X_grid, regressor.predict(X_grid), color='r')
plt.show()