| 219 | } |
| 220 | |
| 221 | fn kmeans_centroids(data: &[&[f32]], dim: usize, k: usize, max_iter: usize) -> Vec<Vec<f32>> { |
| 222 | let n = data.len(); |
| 223 | let k = k.min(n); |
| 224 | if k == 0 { |
| 225 | return Vec::new(); |
| 226 | } |
| 227 | |
| 228 | let mut centroids: Vec<Vec<f32>> = vec![data[0].to_vec()]; |
| 229 | let mut min_dists = vec![f32::MAX; n]; |
| 230 | |
| 231 | // Initialize min_dists against the first centroid. |
| 232 | for (i, point) in data.iter().enumerate() { |
| 233 | let d = distance(point, ¢roids[0], DistanceMetric::L2); |
| 234 | if d < min_dists[i] { |
| 235 | min_dists[i] = d; |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | let mut rng = crate::hnsw::Xorshift64::new(0xC0FF_EEDE_ADBE_EF42); |
| 240 | for _ in 1..k { |
| 241 | let total: f64 = min_dists.iter().map(|&d| d as f64).sum(); |
| 242 | let next_idx = if total < f64::EPSILON { |
| 243 | 0 |
| 244 | } else { |
| 245 | let target = rng.next_f64() * total; |
| 246 | let mut acc = 0.0f64; |
| 247 | let mut chosen = n - 1; |
| 248 | for (i, &d) in min_dists.iter().enumerate() { |
| 249 | acc += d as f64; |
| 250 | if acc >= target { |
| 251 | chosen = i; |
| 252 | break; |
| 253 | } |
| 254 | } |
| 255 | chosen |
| 256 | }; |
| 257 | centroids.push(data[next_idx].to_vec()); |
| 258 | let last = centroids.last().expect("just pushed"); |
| 259 | for (i, point) in data.iter().enumerate() { |
| 260 | let d = distance(point, last, DistanceMetric::L2); |
| 261 | if d < min_dists[i] { |
| 262 | min_dists[i] = d; |
| 263 | } |
| 264 | } |
| 265 | } |
| 266 | |
| 267 | let mut assignments = vec![0usize; n]; |
| 268 | for _ in 0..max_iter { |
| 269 | let mut changed = false; |
| 270 | for (i, point) in data.iter().enumerate() { |
| 271 | let mut best = 0; |
| 272 | let mut best_d = f32::MAX; |
| 273 | for (c, centroid) in centroids.iter().enumerate() { |
| 274 | let d = distance(point, centroid, DistanceMetric::L2); |
| 275 | if d < best_d { |
| 276 | best_d = d; |
| 277 | best = c; |
| 278 | } |