Lloyd's k-means clustering. Returns `k` centroids of length `sub_dim`, initialized via k-means++.
(
points: &[Vec<f32>],
sub_dim: usize,
k: usize,
iters: usize,
seed: u64,
)
| 14 | /// |
| 15 | /// Returns `k` centroids of length `sub_dim`, initialized via k-means++. |
| 16 | pub fn lloyd( |
| 17 | points: &[Vec<f32>], |
| 18 | sub_dim: usize, |
| 19 | k: usize, |
| 20 | iters: usize, |
| 21 | seed: u64, |
| 22 | ) -> Vec<Vec<f32>> { |
| 23 | let n = points.len(); |
| 24 | if n == 0 || k == 0 { |
| 25 | return Vec::new(); |
| 26 | } |
| 27 | let k = k.min(n); |
| 28 | |
| 29 | let mut rng = Xorshift64::new(seed.wrapping_add(0x9E3779B97F4A7C15)); |
| 30 | let mut centroids: Vec<Vec<f32>> = Vec::with_capacity(k); |
| 31 | centroids.push(points[0].clone()); |
| 32 | |
| 33 | let mut min_dists = vec![f32::MAX; n]; |
| 34 | for (i, p) in points.iter().enumerate() { |
| 35 | min_dists[i] = l2_sq(p, ¢roids[0]); |
| 36 | } |
| 37 | |
| 38 | for _ in 1..k { |
| 39 | let total: f64 = min_dists.iter().map(|&d| d as f64).sum(); |
| 40 | let chosen = if total < f64::EPSILON { |
| 41 | 0usize |
| 42 | } else { |
| 43 | let target = { |
| 44 | let u = (rng.next_u64() >> 11) as f64 / (1u64 << 53) as f64; |
| 45 | u * total |
| 46 | }; |
| 47 | let mut acc = 0.0f64; |
| 48 | let mut idx = n - 1; |
| 49 | for (i, &d) in min_dists.iter().enumerate() { |
| 50 | acc += d as f64; |
| 51 | if acc >= target { |
| 52 | idx = i; |
| 53 | break; |
| 54 | } |
| 55 | } |
| 56 | idx |
| 57 | }; |
| 58 | let new_c = points[chosen].clone(); |
| 59 | for (i, p) in points.iter().enumerate() { |
| 60 | let d = l2_sq(p, &new_c); |
| 61 | if d < min_dists[i] { |
| 62 | min_dists[i] = d; |
| 63 | } |
| 64 | } |
| 65 | centroids.push(new_c); |
| 66 | } |
| 67 | |
| 68 | let mut assignments = vec![0usize; n]; |
| 69 | for _ in 0..iters { |
| 70 | let mut changed = false; |
| 71 | for (i, p) in points.iter().enumerate() { |
| 72 | let best = (0..k) |
| 73 | .min_by(|&a, &b| { |