File size: 1,297 Bytes
760a88c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
from matplotlib import pyplot as plt
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
def calculate_wcss(data):
wcss = []
for i in range(1, 11):
kmeans = KMeans(n_clusters=i, init='k-means++', max_iter=300, n_init=10, random_state=0)
kmeans.fit(data)
wcss.append(kmeans.inertia_)
return wcss
def calculate_silhouette_scores(data):
scores = []
range_values = range(2, 11)
for i in range_values:
kmeans = KMeans(n_clusters=i, init='k-means++', max_iter=300, n_init=10, random_state=0)
kmeans.fit(data)
score = silhouette_score(data, kmeans.labels_, metric='euclidean')
scores.append(score)
return scores
def plot_elbow(wcss):
plt.plot(range(1, 11), wcss)
plt.title('Elbow Method')
plt.xlabel('Number of clusters')
plt.ylabel('WCSS')
plt.show()
def get_optimal_clusters_silhouette(scores):
optimal_clusters = scores.index(max(scores)) + 2 # +2 because range_values starts from 2
print(f"Optimal number of clusters: {optimal_clusters}")
return optimal_clusters
def fit_kmeans(data, n_clusters):
kmeans = KMeans(n_clusters=n_clusters, random_state=0)
clusters = kmeans.fit_predict(data)
data['cluster'] = clusters
return kmeans, data |