Table of Contents
Fetching ...

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.

Log-Time K-Means Clustering for 1D Data: Novel Approaches with Proof and Implementation

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 after sorting, and a binary-search-based 2-cluster method with 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 over scikit-learn) and enabling fast applications such as LLM quantization (up to 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 -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 -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 -cluster algorithm achieving complexity for greedy -means++ initialization and for Lloyd's algorithm, where is the number of greedy -means++ local trials, and is the number of Lloyd's algorithm iterations, and (2) a binary search-based two-cluster algorithm, achieving 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 -means clustering, delivering efficient and sound algorithms implemented in a JIT-optimized open-source Python library.

Paper Structure

This paper contains 39 sections, 13 equations, 2 figures, 1 table.

Figures (2)

  • Figure 1: Squared error comparison of the proposed two-cluster algorithm and $k$-cluster algorithms in flash1dkmeans against scikit-learn on real and synthetic datasets. Lower is better.
  • Figure 2: Runtime comparison of the proposed two-cluster algorithm and $k$-cluster algorithms in flash1dkmeans against scikit-learn on datasets of varying sizes. Lower is better.