|
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 |
|
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 |
