kmeans-seeding 0.2.0

Fast k-means++ seeding algorithms with C++ implementation and Python bindings

kmeans-seeding: Fast k-means++ Initialization Algorithms

Fast, state-of-the-art k-means initialization algorithms implemented in C++ with Python bindings.

Features

🚀 Fast: C++ implementation with OpenMP parallelization 🎯 Accurate: State-of-the-art algorithms with theoretical guarantees 🔌 Compatible: Drop-in replacement for sklearn's k-means++ initialization 📦 Easy to use: Simple Python API, works with NumPy arrays 🛠️ Flexible: Multiple algorithms to choose from

Algorithms Included

1. RS-k-means++ (Rejection Sampling) - Our contribution

   • Fast approximate D² sampling using rejection sampling
   • Supports FAISS for approximate nearest neighbors
   • Best for large datasets (n > 10,000)

2. AFK-MC² (Adaptive Fast k-MC²)

   • MCMC-based sampling without computing all distances
   • Good balance of speed and quality

3. Fast-LSH k-means++ (Google 2020)

   • Tree embedding with LSH for fast sampling
   • Excellent for high-dimensional data

4. Standard k-means++

   • Classic D² sampling algorithm
   • Baseline for comparison

Installation

From PyPI (recommended)

pip install kmeans-seeding

With FAISS support (recommended for large datasets)

# CPU version
conda install -c pytorch faiss-cpu
pip install kmeans-seeding

# GPU version
conda install -c pytorch faiss-gpu
pip install kmeans-seeding

From source

git clone https://github.com/pcshah2004/kmeans-seeding.git
cd kmeans-seeding
pip install -e .

Quick Start

from kmeans_seeding import rejection_sampling
from sklearn.cluster import KMeans
import numpy as np

# Generate sample data
X = np.random.randn(10000, 50)

# Get initial centers using RS-k-means++
centers = rejection_sampling(X, n_clusters=100, index_type='LSH')

# Use with sklearn
kmeans = KMeans(n_clusters=100, init=centers, n_init=1)
kmeans.fit(X)

Usage Examples

RS-k-means++ (Rejection Sampling)

from kmeans_seeding import rejection_sampling

centers = rejection_sampling(
    X,
    n_clusters=100,
    max_iter=50,           # Max rejection sampling iterations
    index_type='LSH',      # FAISS index type: 'Flat', 'LSH', 'IVFFlat', 'HNSW'
    random_state=42
)

AFK-MC² (MCMC Sampling)

from kmeans_seeding import afkmc2

centers = afkmc2(
    X,
    n_clusters=100,
    chain_length=200,      # Markov chain length
    random_state=42
)

Fast-LSH k-means++

from kmeans_seeding import fast_lsh

centers = fast_lsh(
    X,
    n_clusters=100,
    n_trees=4,             # Number of trees for embedding
    random_state=42
)

Standard k-means++

from kmeans_seeding import kmeanspp

centers = kmeanspp(X, n_clusters=100, random_state=42)

Benchmarks

Performance comparison on various datasets:

Dataset	n	d	Algorithm	Time (s)	Cost Ratio*
MNIST	60K	784	k-means++	45.2	1.00
			RS-k-means++ (LSH)	2.1	1.02
			AFK-MC²	8.3	1.05
CIFAR-10	50K	512	k-means++	38.7	1.00
			RS-k-means++ (LSH)	1.8	1.01
			AFK-MC²	6.9	1.04

*Cost ratio: Final k-means cost compared to standard k-means++

Documentation

Full documentation available at: https://kmeans-seeding.readthedocs.io

• Installation Guide
• API Reference
• Algorithm Details
• Benchmarks

Requirements

• Python 3.9+
• NumPy >= 1.20.0
• (Optional) FAISS >= 1.7.0 for fast approximate nearest neighbors
• (Optional) scikit-learn for full k-means clustering

Citation

If you use this library in your research, please cite:

@article{shah2025rejection,
  title={A New Rejection Sampling Approach to k-means++ With Improved Trade-Offs},
  author={Shah, Poojan and Agrawal, Shashwat and Jaiswal, Ragesh},
  journal={arXiv preprint arXiv:2502.02085},
  year={2025}
}

For AFK-MC²:

@inproceedings{bachem2016approximate,
  title={Approximate k-means++ in sublinear time},
  author={Bachem, Olivier and Lucic, Mario and Hassani, Hamed and Krause, Andreas},
  booktitle={AAAI Conference on Artificial Intelligence},
  year={2016}
}

For Fast-LSH k-means++:

@inproceedings{cohen2020fast,
  title={Fast and accurate k-means++ via rejection sampling},
  author={Cohen-Addad, Vincent and Lattanzi, Silvio and Mitrovi{\'c}, Slobodan and Norouzi-Fard, Ashkan and Parotsidis, Nikos and Tarnawski, Jakub},
  booktitle={Advances in Neural Information Processing Systems},
  year={2020}
}

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• FAISS library by Facebook AI Research
• scikit-learn for the k-means clustering API design
• Research supported by the Department of Computer Science, IIT Delhi

Contact

• Poojan Shah: [e-mail]
• Issues: https://github.com/pcshah2004/kmeans-seeding/issues
• Discussions: https://github.com/pcshah2004/kmeans-seeding/discussions