Recompute centroids as the mean of their assigned vectors. Empty clusters are re-seeded to the input vector farthest from all live centroids — this guarantees every centroid covers some part of the input space and prevents k-means from collapsing into fewer-than-k effective clusters.
(
vectors: &[Vec<f32>],
assignments: &[usize],
num_centroids: usize,
dim: usize,
prev_centroids: &[Vec<f32>],
)
| 59 | /// guarantees every centroid covers some part of the input space and prevents |
| 60 | /// k-means from collapsing into fewer-than-k effective clusters. |
| 61 | fn recompute( |
| 62 | vectors: &[Vec<f32>], |
| 63 | assignments: &[usize], |
| 64 | num_centroids: usize, |
| 65 | dim: usize, |
| 66 | prev_centroids: &[Vec<f32>], |
| 67 | ) -> Vec<Vec<f32>> { |
| 68 | let mut sums = vec![vec![0.0f32; dim]; num_centroids]; |
| 69 | let mut counts = vec![0usize; num_centroids]; |
| 70 | |
| 71 | for (v, &c) in vectors.iter().zip(assignments.iter()) { |
| 72 | for (s, x) in sums[c].iter_mut().zip(v.iter()) { |
| 73 | *s += x; |
| 74 | } |
| 75 | counts[c] += 1; |
| 76 | } |
| 77 | |
| 78 | // First pass: average populated clusters in place. |
| 79 | for (s, &n) in sums.iter_mut().zip(counts.iter()) { |
| 80 | if n > 0 { |
| 81 | s.iter_mut().for_each(|x| *x /= n as f32); |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | // Second pass: re-seed empty clusters from vectors farthest from any live |
| 86 | // centroid. Snapshot the live set first so each empty slot is filled |
| 87 | // deterministically and subsequent re-seeds in the same call see the |
| 88 | // updated pool. |
| 89 | for c_idx in 0..num_centroids { |
| 90 | if counts[c_idx] != 0 { |
| 91 | continue; |
| 92 | } |
| 93 | let live: Vec<Vec<f32>> = counts |
| 94 | .iter() |
| 95 | .enumerate() |
| 96 | .filter(|(i, cnt)| *i != c_idx && **cnt > 0) |
| 97 | .map(|(i, _)| sums[i].clone()) |
| 98 | .collect(); |
| 99 | let seed_pool: &[Vec<f32>] = if live.is_empty() { |
| 100 | prev_centroids |
| 101 | } else { |
| 102 | &live |
| 103 | }; |
| 104 | let farthest = vectors |
| 105 | .iter() |
| 106 | .map(|v| (v, min_dist_to_centroids(v, seed_pool))) |
| 107 | .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) |
| 108 | .map(|(v, _)| v.clone()); |
| 109 | if let Some(v) = farthest { |
| 110 | sums[c_idx] = v; |
| 111 | counts[c_idx] = 1; // mark live so later empty slots see it as a seed. |
| 112 | } else if c_idx < prev_centroids.len() { |
| 113 | sums[c_idx] = prev_centroids[c_idx].clone(); |
| 114 | counts[c_idx] = 1; |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | sums |
no test coverage detected