gistfile1.txt

km = KMeans(n_clusters = 5)
y_predicted = km.fit_predict(salary_story[['Total Salary Paid', 'Home Price']])
salary_story['Cluster'] = y_predicted

gistfile2.txt

plt.figure(figsize = (10,5))

df1 = salary_story[salary_story.Cluster == 0]
df2 = salary_story[salary_story.Cluster == 1]
df3 = salary_story[salary_story.Cluster == 2]
df4 = salary_story[salary_story.Cluster == 3]
df5 = salary_story[salary_story.Cluster == 4]

plt.scatter(df1['Total Salary Paid'], df1['Home Price'], color = 'green')
plt.scatter(df2['Total Salary Paid'], df2['Home Price'], color = 'red')
plt.scatter(df3['Total Salary Paid'], df3['Home Price'], color = 'yellow')
plt.scatter(df4['Total Salary Paid'], df4['Home Price'], color='blue')
plt.scatter(df5['Total Salary Paid'], df5['Home Price'], color='black')


plt.scatter(km.cluster_centers_[:, 0], km.cluster_centers_[:, 1], color='purple', marker='*', label='centroid')

plt.xlabel('Total Salary Paid')
plt.ylabel('Home Price')
plt.legend()