k-means_exercise_1.py

from sklearn.cluster import KMeans
import numpy as np

# Generate data to be clustered
X = np.array([[1, 2], [1, 4], [1, 0],
              [10, 2], [10, 4], [10, 0]])

# Init K-menas model and clustering
kmeans = KMeans(n_clusters=2).fit(X)

# Get labels for each cluster
kmeans.labels_
# Output:
# array([1, 1, 1, 0, 0, 0], dtype=int32)

# Predict new data point
kmeans.predict([[0, 0], [12, 3]])
# Output:
# array([1, 0], dtype=int32)

# Get centroid of clusters
kmeans.cluster_centers_
# Outcome:
# array([[10.,  2.],
#        [ 1.,  2.]])