API Reference

Main Functions

All algorithms return cluster centers as NumPy arrays of shape (n_clusters, n_features) that can be used directly with scikit-learn’s KMeans.

rskmeans

kmeans_seeding.rskmeans(X, n_clusters, *, max_iter=50, index_type='LSH', random_state=None)[source]

RS-k-means++ initialization using rejection sampling.

Fast k-means++ seeding that uses rejection sampling with approximate nearest neighbor queries for efficient D² sampling.

Parameters:

X (array-like of shape (n_samples, n_features)) – Training data.
n_clusters (int) – Number of clusters.
max_iter (int, default=50) – Maximum number of rejection sampling iterations per center. Higher values improve quality but take longer.
index_type ({'Flat', 'LSH', 'IVFFlat', 'HNSW', 'FastLSH'}, default='LSH') – Approximate nearest neighbor index type: - ‘Flat’: Exact search (slowest, most accurate) - ‘LSH’: FAISS LSH (fast, ~90-95% accuracy) - ‘IVFFlat’: Inverted file index (fast, ~99% accuracy) - ‘HNSW’: Hierarchical NSW (very fast, ~95-99% accuracy) - ‘FastLSH’: DHHash-based Fast LSH (very fast, ~90-95% accuracy)
random_state (int, optional) – Random seed for reproducibility.

Returns:

centers – Initial cluster centers selected using RS-k-means++.

Return type:

ndarray of shape (n_clusters, n_features)

References

Shah, P., Agrawal, S., & Jaiswal, R. (2025). A New Rejection Sampling Approach to k-means++ With Improved Trade-Offs. arXiv preprint arXiv:2502.02085.

Examples

>>> from kmeans_seeding import rskmeans
>>> from sklearn.cluster import KMeans
>>> import numpy as np
>>>
>>> X = np.random.randn(10000, 50)
>>> centers = rskmeans(X, n_clusters=100, index_type='LSH')
>>> kmeans = KMeans(n_clusters=100, init=centers, n_init=1)
>>> kmeans.fit(X)

kmeanspp

kmeans_seeding.kmeanspp(X, n_clusters, *, random_state=None)[source]

Standard k-means++ initialization.

Parameters:

X (array-like of shape (n_samples, n_features)) – Training data.
n_clusters (int) – Number of clusters.
random_state (int, optional) – Random seed for reproducibility.

Returns:

centers – Initial cluster centers selected using k-means++.

Return type:

ndarray of shape (n_clusters, n_features)

References

Arthur, D., & Vassilvitskii, S. (2007). k-means++: The advantages of careful seeding.

Examples

>>> from kmeans_seeding import kmeanspp
>>> from sklearn.cluster import KMeans
>>> import numpy as np
>>>
>>> X = np.random.randn(1000, 10)
>>> centers = kmeanspp(X, n_clusters=5)
>>> kmeans = KMeans(n_clusters=5, init=centers, n_init=1)
>>> kmeans.fit(X)

afkmc2

kmeans_seeding.afkmc2(X, n_clusters, *, chain_length=200, index_type='Flat', random_state=None)[source]

AFK-MC² initialization using MCMC sampling.

Adaptive Fast k-MC² uses Markov Chain Monte Carlo to sample centers according to the D² distribution without computing all distances.

Parameters:

X (array-like of shape (n_samples, n_features)) – Training data.
n_clusters (int) – Number of clusters.
chain_length (int, default=200) – Length of the Markov chain per center. Longer chains give better quality but take more time.
index_type ({'Flat', 'LSH', 'HNSW'}, default='Flat') – FAISS index type for label assignment (not used in sampling).
random_state (int, optional) – Random seed for reproducibility.

Returns:

centers – Initial cluster centers selected using AFK-MC².

Return type:

ndarray of shape (n_clusters, n_features)

References

Bachem, O., Lucic, M., Hassani, H., & Krause, A. (2016). Approximate k-means++ in sublinear time. AAAI Conference on Artificial Intelligence.

Examples

>>> from kmeans_seeding import afkmc2
>>> from sklearn.cluster import KMeans
>>> import numpy as np
>>>
>>> X = np.random.randn(10000, 50)
>>> centers = afkmc2(X, n_clusters=100, chain_length=200)
>>> kmeans = KMeans(n_clusters=100, init=centers, n_init=1)
>>> kmeans.fit(X)

multitree_lsh

kmeans_seeding.multitree_lsh(X, n_clusters, *, n_trees=4, scaling_factor=1.0, n_greedy_samples=1, index_type='Flat', random_state=None)[source]

Fast k-means++ using tree embedding (Google 2020).

Uses tree embedding and integer casting for fast D² sampling.

Parameters:

X (array-like of shape (n_samples, n_features)) – Training data.
n_clusters (int) – Number of clusters.
n_trees (int, default=4) – Number of trees for embedding.
scaling_factor (float, default=1.0) – Scaling factor for integer casting.
n_greedy_samples (int, default=1) – Number of greedy samples per center.
index_type (str, default='Flat') – FAISS index type for label assignment.
random_state (int, optional) – Random seed for reproducibility.

Returns:

centers – Initial cluster centers.

Return type:

ndarray of shape (n_clusters, n_features)

References

Cohen-Addad, V., Lattanzi, S., Mitrović, S., Norouzi-Fard, A., Parotsidis, N., & Tarnawski, J. (2020). Fast and accurate k-means++ via rejection sampling. NeurIPS 2020.

Examples

>>> from kmeans_seeding import multitree_lsh
>>> from sklearn.cluster import KMeans
>>> import numpy as np
>>>
>>> X = np.random.randn(10000, 50)
>>> centers = multitree_lsh(X, n_clusters=100)
>>> kmeans = KMeans(n_clusters=100, init=centers, n_init=1)
>>> kmeans.fit(X)

Aliases

rejection_sampling

Alias for rskmeans().

from kmeans_seeding import rejection_sampling

# Same as rskmeans
centers = rejection_sampling(X, n_clusters=100)

fast_lsh

Alias for multitree_lsh().

from kmeans_seeding import fast_lsh

# Same as multitree_lsh
centers = fast_lsh(X, n_clusters=100)

Complete Example

from kmeans_seeding import (
    kmeanspp,
    rskmeans,
    afkmc2,
    multitree_lsh,
    rejection_sampling,  # alias for rskmeans
    fast_lsh,            # alias for multitree_lsh
)
from sklearn.cluster import KMeans
import numpy as np

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

# All algorithms have the same basic interface
centers_kpp = kmeanspp(X, n_clusters=k, random_state=42)
centers_rs = rskmeans(X, n_clusters=k, random_state=42)
centers_afk = afkmc2(X, n_clusters=k, random_state=42)
centers_lsh = multitree_lsh(X, n_clusters=k, random_state=42)

# Use with scikit-learn
kmeans = KMeans(n_clusters=k, init=centers_rs, n_init=1)
labels = kmeans.fit_predict(X)

Parameter Summary

Common Parameters

All functions share these parameters:

Parameter	Type	Required	Description
`X`	array-like	Yes	Training data, shape (n_samples, n_features)
`n_clusters`	int	Yes	Number of clusters to initialize
`random_state`	int	No	Random seed for reproducibility

Algorithm-Specific Parameters

rskmeans:

`max_iter`	int	Max rejection sampling iterations (default: 50)
`index_type`	str	ANN index: ‘Flat’, ‘LSH’, ‘IVFFlat’, ‘HNSW’, ‘FastLSH’, ‘GoogleLSH’

afkmc2:

`chain_length`	int	Markov chain length per center (default: 200)
`index_type`	str	FAISS index for label assignment (default: ‘Flat’)

multitree_lsh:

`n_trees`	int	Number of random projection trees (default: 4)
`scaling_factor`	float	Integer casting precision (default: 1.0)
`n_greedy_samples`	int	Greedy samples per center (default: 1)
`index_type`	str	FAISS index for label assignment (default: ‘Flat’)

Return Values

All functions return:

centers : ndarray of shape (n_clusters, n_features)
    The initialized cluster centers.

These can be used directly with scikit-learn:

from sklearn.cluster import KMeans

kmeans = KMeans(n_clusters=k, init=centers, n_init=1)
labels = kmeans.fit_predict(X)

Exceptions

Common exceptions raised by all algorithms:

ValueError:

n_samples < n_clusters
Invalid parameter values
Dimension mismatch

RuntimeError (rskmeans only):

FAISS index requested but FAISS not available
Invalid index type

Example:

from kmeans_seeding import rskmeans
import numpy as np

X = np.random.randn(100, 10)

try:
    # This will fail: FAISS not installed
    centers = rskmeans(X, n_clusters=10, index_type='LSH')
except RuntimeError as e:
    print(f"Error: {e}")
    # Use FastLSH instead (works without FAISS)
    centers = rskmeans(X, n_clusters=10, index_type='FastLSH')

Type Hints

For type checkers and IDEs:

from typing import Optional
import numpy as np
from numpy.typing import ArrayLike, NDArray

def kmeanspp(
    X: ArrayLike,
    n_clusters: int,
    *,
    random_state: Optional[int] = None
) -> NDArray[np.float64]:
    ...

def rskmeans(
    X: ArrayLike,
    n_clusters: int,
    *,
    max_iter: int = 50,
    index_type: str = 'LSH',
    random_state: Optional[int] = None
) -> NDArray[np.float64]:
    ...

def afkmc2(
    X: ArrayLike,
    n_clusters: int,
    *,
    chain_length: int = 200,
    index_type: str = 'Flat',
    random_state: Optional[int] = None
) -> NDArray[np.float64]:
    ...

def multitree_lsh(
    X: ArrayLike,
    n_clusters: int,
    *,
    n_trees: int = 4,
    scaling_factor: float = 1.0,
    n_greedy_samples: int = 1,
    index_type: str = 'Flat',
    random_state: Optional[int] = None
) -> NDArray[np.float64]:
    ...

API Reference

Main Functions

rskmeans

kmeanspp

afkmc2

multitree_lsh

Aliases

rejection_sampling

fast_lsh

Complete Example

Parameter Summary

Common Parameters

Algorithm-Specific Parameters

Return Values

Exceptions

Type Hints

See Also