# K-Means: Find the Groups
# Core implementation example — ackleywill.com

import numpy as np

X = np.array([[1., 1.], [1., 2.], [4., 4.], [5., 4.]])
centers = X[[0, 2]].copy()

for _ in range(20):
    # Assign each point to its nearest center.
    distance2 = ((X[:, None, :] - centers[None, :, :]) ** 2).sum(axis=2)
    labels = distance2.argmin(axis=1)
    # Keep an empty cluster's previous center.
    updated = np.array([
        X[labels == j].mean(axis=0) if np.any(labels == j) else centers[j]
        for j in range(len(centers))
    ])
    if np.allclose(updated, centers):
        break
    centers = updated

print(centers)