urschrei/ckmeans

Optimal univariate k-means clustering using dynamic programming

Ckmeans

use ckmeans::ckmeans;

let input = vec![
    1.0, 12.0, 13.0, 14.0, 15.0, 16.0, 2.0,
    2.0, 3.0, 5.0, 7.0, 1.0, 2.0, 5.0, 7.0,
    1.0, 5.0, 82.0, 1.0, 1.3, 1.1, 78.0,
];
let expected = vec![
    vec![
        1.0, 1.0, 1.0, 1.0, 1.1, 1.3, 2.0, 2.0,
        2.0, 3.0, 5.0, 5.0, 5.0, 7.0, 7.0,
    ],
    vec![12.0, 13.0, 14.0, 15.0, 16.0],
    vec![78.0, 82.0],
];
let result = ckmeans(&input, 3).unwrap();
assert_eq!(result, expected);

Ckmeans clustering is an improvement on 1-dimensional (univariate) heuristic-based clustering approaches such as Jenks. The algorithm was developed by Haizhou Wang and Mingzhou Song (2011) as a dynamic programming approach to the problem of clustering numeric data into groups with the least within-group sum-of-squared-deviations.

Minimizing the difference within groups – what Wang & Song refer to as withinss, or within sum-of-squares – means that groups are optimally homogenous within and the data is split into representative groups. This is very useful for visualization, where one may wish to represent a continuous variable in discrete colour or style groups. This function can provide groups that emphasize differences between data.

Being a dynamic approach, this algorithm is based on two matrices that store incrementally-computed values for squared deviations and backtracking indexes.

Unlike the original implementation, this implementation does not include any code to automatically determine the optimal number of clusters: this information needs to be explicitly provided. It does provide the roundbreaks method to aid labelling, however.

FFI

A C-compatible FFI implementation is available, along with libraries for major platforms. See the header file and a basic C example in the examples folder. The FFI functions have been verified not to leak memory (see comment in example).

WASM

A WASM module is also available, giving access to both ckmeans and roundbreaks. Generate the module using wasm-bindgen and the appropriate target, or use the NPM package.

Implementation

This is a port (including documentation) of David Schnurr's package https://github.com/schnerd/ckmeans, incorporating some improvements from Bill Mill's Python + Numpy implementation at https://github.com/llimllib/ckmeans.

Performance

On an M2 Pro, to produce 7 classes:

1. 110k uniformly-distributed i32 values between 0 and 250: ~12 ms
2. 110k normally-distributed f64 values with a mean of 3.0 and a standard deviation of 1.0: 38 ms

Complexity

$O(kn)$. Other approaches such as Hilferink's CalcNaturalBreaks or k-means have comparable complexity, but do not guarantee optimality. In practice, they require many rounds to approach an optimal result, so in practice they're slower.

Note

Wang and Song (2011) state that the algorithm runs in $O(k^2n)$ in their introduction. However, they have since updated their dynamic programming algorithm (see August 2016 note here) which reduces the complexity to linear time. This approach has been used in the extant implementations listed above, and reproduced here.

Possible Improvements

Perf

The "matrices" are nested vectors and thus don't have optimal memory layout. In addition, we're not trying to leverage any of the fast linear algebra libraries that might be available if we used e.g. ndarray.

Tests

Perhaps some property-based tests.

References

1. Wang, H., & Song, M. (2011). Ckmeans.1d.dp: Optimal k-means Clustering in One Dimension by Dynamic Programming. The R Journal, 3(2), 29.
2. https://observablehq.com/@visionscarto/natural-breaks