import numpy as np
from sklearn.neighbors import KNeighborsRegressor

# Prepare the Data
# X = Square Footage (must be 2D array)
X_train = np.array([[800], [1200], [1500], [1800], [2200], [2500]])

# y = Price in Dollars
y_train = np.array([150000, 230000, 280000, 310000, 400000, 450000])

# Initialize the K-NN Regressor
# Look at the 3 nearest neighbors (k=3)
knn_reg = KNeighborsRegressor(n_neighbors=3)

# Train the model
knn_reg.fit(X_train, y_train)

# Predict the price for a house with 2000 sq ft
mystery_house = np.array([[2000]])
predicted_price = knn_reg.predict(mystery_house)

print(f"Predicted price for a 2000 sq ft house N=3: ${predicted_price[0]:,.2f}")

# Initialize the K-NN Regressor
# Look at the 2 nearest neighbors (k=2)
knn_reg = KNeighborsRegressor(n_neighbors=2)

# Train the model
knn_reg.fit(X_train, y_train)

# Predict the price for a house with 2000 sq ft
mystery_house = np.array([[2000]])
predicted_price = knn_reg.predict(mystery_house)

print(f"Predicted price for a 2000 sq ft house N=2: ${predicted_price[0]:,.2f}")