text_plain_output_1.png image_output_1.png
from sklearn.cluster import KMeans

import matplotlib.pyplot as plt

import numpy as np

import numpy as np # linear algebra

import pandas as pd

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

workdir = '../input/'
basins_data = pd.read_csv(workdir + 'BasinCharacteristics.csv')

features_cols = basins_data.columns[1:4]
target_col = basins_data.columns[0]
X = basins_data[features_cols]
y = basins_data[target_col]

def elbow_plot(data, maxK=10, seed_centroids=None):
    """
        parameters:
        - data: pandas DataFrame (data to be fitted)
        - maxK (default = 10): integer (maximum number of clusters with which to run k-means)
        - seed_centroids (default = None ): float (initial value of centroids for k-means)
    """
    sse = {}
    for k in range(1, maxK):
        print('k: ', k)
        if seed_centroids is not None:
            seeds = seed_centroids.head(k)
            kmeans = KMeans(n_clusters=k, max_iter=500, n_init=100, random_state=0, init=np.reshape(seeds, (k, 1))).fit(data)
            data['clusters'] = kmeans.labels_
        else:
            kmeans = KMeans(n_clusters=k, max_iter=300, n_init=100, random_state=0).fit(data)
            data['clusters'] = kmeans.labels_
        sse[k] = kmeans.inertia_
    plt.figure()
    plt.plot(list(sse.keys()), list(sse.values()))
    plt.show()
    return
elbow_plot(X, maxK=10)