K-Means++

Fast KMeans++ initialization.
A scikit-learn-compatible KMeans implementation with fast SIMD distance computations and parallel centroid initialization.

Documentation · PyPI

fastkmeanspp is a Python package that implements a KMeans clone from scikit-learn with a faster KMeans++ centroid initialization. It is designed to be a drop-in replacement for scikit-learn's KMeans when initialization is the bottleneck.

✨ Features

• Fast KMeans++ initialization: Uses optimized squared-distance computations while selecting candidate centroids.
• Portable SIMD kernels: Uses Google Highway to compile vectorized kernels for supported CPU targets and select the best implementation at runtime.
• Parallel initialization: Computes distance rows in parallel with Highway's thread pool.
• scikit-learn compatibility: Provides familiar fit, predict, labels_, cluster_centers_, and inertia_ interfaces.
• SIMD Lloyd updates: Uses the same native kernels for nearest-centroid assignment, cluster accumulation, and centroid updates.

How Google Highway fits in

Google Highway is a C++ library for portable SIMD programming. The distance and clustering kernels are written once with Highway's vector operations. Highway then builds target-specific versions for the available instruction sets, such as SSE, AVX2, and AVX-512 on x86 CPUs. A small runtime dispatch layer selects the strongest target supported by the current processor and keeps a scalar fallback for portability.

The distance kernel loads several feature values at a time, subtracts the corresponding centroid values, and uses fused multiply-add operations to build the squared distance. Feature dimensions that do not fill a complete vector are handled by a short scalar tail.

During KMeans++ initialization, Highway's thread pool divides independent data rows between workers. During Lloyd updates, each worker finds the closest centroid and accumulates feature sums and point counts. The native update then turns those sums into new centroid coordinates without changing the estimator interface.

🚀 Installation

python -m pip install fastkmeanspp

🔧 Usage

import numpy as np
from fastkmeanspp import KMeans

X = np.array([[0.0, 0.0], [0.1, 0.2], [4.0, 4.0], [4.2, 3.9]])
model = KMeans(n_clusters=2, random_state=42)
model.fit(X)

print(model.labels_)
print(model.cluster_centers_)

Set n_jobs=1 for serial centroid initialization or n_jobs=-1 to use all available threads.