Log-Time K-Means Clustering for 1D Data: Novel Approaches with Proof and Implementation
Jake Hyun
TL;DR
This work addresses the computational challenge of 1D k-means clustering by exploiting the natural order of data. It introduces two log-time approaches: a k-cluster method that combines greedy k-means++ initialization with Lloyd iterations, achieving a complexity of $O(l k^2 log n) + O(i k log n)$ after sorting, and a binary-search-based 2-cluster method with $O(log n)$ runtime that deterministically converges to a Lloyd's local minimum. Through prefix sums and careful boundary management, both algorithms preserve clustering quality while delivering substantial speedups (up to $4500\times$ over scikit-learn) and enabling fast applications such as LLM quantization (up to $300\times$ speedups). An open-source, JIT-optimized Python library, flash1dkmeans, implements these methods, bridging theory and practical deployment for latency-critical 1D clustering tasks.
Abstract
Clustering is a key task in machine learning, with $k$-means being widely used for its simplicity and effectiveness. While 1D clustering is common, existing methods often fail to exploit the structure of 1D data, leading to inefficiencies. This thesis introduces optimized algorithms for $k$-means++ initialization and Lloyd's algorithm, leveraging sorted data, prefix sums, and binary search for improved computational performance. The main contributions are: (1) an optimized $k$-cluster algorithm achieving $O(l \cdot k^2 \cdot \log n)$ complexity for greedy $k$-means++ initialization and $O(i \cdot k \cdot \log n)$ for Lloyd's algorithm, where $l$ is the number of greedy $k$-means++ local trials, and $i$ is the number of Lloyd's algorithm iterations, and (2) a binary search-based two-cluster algorithm, achieving $O(\log n)$ runtime with deterministic convergence to a Lloyd's algorithm local minimum. Benchmarks demonstrate over a 4500x speedup compared to scikit-learn for large datasets while maintaining clustering quality measured by within-cluster sum of squares (WCSS). Additionally, the algorithms achieve a 300x speedup in an LLM quantization task, highlighting their utility in emerging applications. This thesis bridges theory and practice for 1D $k$-means clustering, delivering efficient and sound algorithms implemented in a JIT-optimized open-source Python library.
