Skip to content

K-means

Algorithm

K-means partitions observations into \(k\) clusters, minimizing the sum of squared distances from each observation \(x\) to its cluster centroid \(\mu_j\):

\[ \underset{C_1,\ldots,C_k}{\operatorname{minimize}} \sum_{j=1}^{k}\sum_{x\in C_j}\lVert x-\mu_j\rVert_2^2 \]

Each iteration assigns observations to the nearest centroid and recomputes the centroids from those assignments. Initialization affects which local solution is found.

Inputs and limitations

  • Pass a non-empty rectangular list of finite numeric observations and an integer 1 <= k <= len(samples).
  • Distances are Euclidean; scale features first when their units differ.
  • Choose centroids_init_method="random" or "kmeans++".
  • Duplicate locations can produce fewer non-empty clusters than k, with a warning log. All k centers retain valid coordinates.
  • Set random_seed to an integer for reproducible fits.
  • predict accepts one sample; fit_predict returns one label per training observation.

Minimal example

from toyml.clustering import KMeans

samples = [[0.0], [0.1], [4.0], [4.1]]
model = KMeans(k=2, random_seed=42)
print(model.fit_predict(samples))
print(model.predict([0.2]))

API

toyml.clustering.kmeans.KMeans dataclass 

KMeans(k: int, max_iter: int = 500, tol: float = 1e-05, centroids_init_method: Literal['random', 'kmeans++'] = 'random', random_seed: int | None = None, distance_metric: Literal['euclidean'] = 'euclidean')

K-means algorithm (with KMeans++ initialization as option).

Examples:

>>> from toyml.clustering import KMeans
>>> dataset = [[1.0, 2.0], [1.0, 4.0], [1.0, 0.0], [10.0, 2.0], [10.0, 4.0], [11.0, 0.0]]
>>> kmeans = KMeans(k=2, random_seed=42).fit(dataset)
>>> kmeans.clusters_
{0: [3, 4, 5], 1: [0, 1, 2]}
>>> kmeans.centroids_
{0: [10.333333333333334, 2.0], 1: [1.0, 2.0]}
>>> kmeans.labels_
[1, 1, 1, 0, 0, 0]
>>> kmeans.predict([0, 1])
1
>>> kmeans.iter_
2

There is a fit_predict method that can be used to fit and predict.

Examples:

>>> from toyml.clustering import KMeans
>>> dataset = [[1, 0], [1, 1], [1, 2], [10, 0], [10, 1], [10, 2]]
>>> KMeans(k=2, random_seed=42).fit_predict(dataset)
[1, 1, 1, 0, 0, 0]
References
  1. Zhou Zhihua
  2. Murphy
Note

Here we just implement the naive K-means algorithm.

See Also

k instance-attribute 

k: int

The number of clusters, specified by user.

max_iter class-attribute instance-attribute 

max_iter: int = 500

The number of iterations the algorithm will run for if it does not converge before that.

tol class-attribute instance-attribute 

tol: float = 1e-05

The tolerance for convergence.

centroids_init_method class-attribute instance-attribute 

centroids_init_method: Literal['random', 'kmeans++'] = 'random'

The method to initialize the centroids.

random_seed class-attribute instance-attribute 

random_seed: int | None = None

Integer seeds reproduce each fit; None keeps an independent random stream.

distance_metric class-attribute instance-attribute 

distance_metric: Literal['euclidean'] = 'euclidean'

The distance metric to use.(For now we only support euclidean).

clusters_ class-attribute instance-attribute 

clusters_: dict[int, list[int]] = field(default_factory=dict, init=False)

The clusters of the dataset.

centroids_ class-attribute instance-attribute 

centroids_: dict[int, list[float]] = field(default_factory=dict, init=False)

The centroids of the clusters.

labels_ class-attribute instance-attribute 

labels_: list[int] = field(default_factory=list, init=False)

The cluster labels of the dataset.

fit 

fit(dataset: list[list[float]]) -> KMeans

Fit the dataset with K-means algorithm.

PARAMETER DESCRIPTION
dataset

the set of data points for clustering

TYPE: list[list[float]]

RETURNS DESCRIPTION
KMeans

self.
Source code in src/toyml/clustering/kmeans.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def fit(self, dataset: list[list[float]]) -> KMeans:
    """Fit the dataset with K-means algorithm.

    Args:
        dataset: the set of data points for clustering

    Returns:
        self.
    """
    self._validate_input(dataset)
    # Replay integer-seeded fits; unseeded instances keep advancing their private stream.
    if self.random_seed is not None:
        self.random_state.seed(self.random_seed)
    self.iter_ = 0
    self.centroids_ = self._get_initial_centroids(dataset)
    for _ in range(self.max_iter):
        self.iter_ += 1
        prev_centroids = self.centroids_
        self._iter_step(dataset)
        if self._is_converged(prev_centroids):
            break
    # Expose assignments to the final centers, including after an iteration limit.
    self.clusters_ = self._get_clusters(dataset)
    self.labels_ = self._predict_dataset_labels(dataset)
    nonempty = sum(bool(cluster) for cluster in self.clusters_.values())
    if nonempty < self.k:
        logger.warning("Found %d non-empty clusters, fewer than requested k=%d", nonempty, self.k)
    return self

fit_predict 

fit_predict(dataset: list[list[float]]) -> list[int]

Fit and predict the cluster label of the dataset.

PARAMETER DESCRIPTION
dataset

the set of data points for clustering

TYPE: list[list[float]]

RETURNS DESCRIPTION
list[int]

Cluster labels of the dataset samples.
Source code in src/toyml/clustering/kmeans.py 
133
134
135
136
137
138
139
140
141
142
def fit_predict(self, dataset: list[list[float]]) -> list[int]:
    """Fit and predict the cluster label of the dataset.

    Args:
        dataset: the set of data points for clustering

    Returns:
        Cluster labels of the dataset samples.
    """
    return self.fit(dataset).labels_.copy()

predict 

predict(point: list[float]) -> int

Predict the label of the point.

PARAMETER DESCRIPTION
point

The data point to predict.

TYPE: list[float]

RETURNS DESCRIPTION
int

The label of the point.
Source code in src/toyml/clustering/kmeans.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
def predict(self, point: list[float]) -> int:
    """Predict the label of the point.

    Args:
        point: The data point to predict.

    Returns:
        The label of the point.

    """
    if len(self.centroids_) == 0:
        msg = "The model is not fitted yet"
        raise ValueError(msg)
    return self._get_point_centroid_label(point, self.centroids_)