K-Means算法

  1. // Macro to implement kmeans for both f64 and f32 without writing everything
  2. // twice or importing the `num` crate
  3. macro_rules! impl_kmeans {
  4.     ($kind: ty, $modname: ident) => {
  5.         // Since we can't overload methods in rust, we have to use namespacing
  6.         pub mod $modname {
  7.             use std::$modname::INFINITY;
  9.             /// computes sum of squared deviation between two identically sized vectors
  10.             /// `x`, and `y`.
  11.             fn distance(x: &[$kind], y: &[$kind]) -> $kind {
  12.                 x.iter()
  13.                     .zip(y.iter())
  14.                     .fold(0.0, |dist, (&xi, &yi)| dist + (xi - yi).powi(2))
  15.             }
  17.             /// Returns a vector containing the indices z<sub>i</sub> in {0, ..., K-1} of
  18.             /// the centroid nearest to each datum.
  19.             fn nearest_centroids(xs: &[Vec<$kind>], centroids: &[Vec<$kind>]) -> Vec<usize> {
  20.                 xs.iter()
  21.                     .map(|xi| {
  22.                         // Find the argmin by folding using a tuple containing the argmin
  23.                         // and the minimum distance.
  24.                         let (argmin, _) = centroids.iter().enumerate().fold(
  25.                             (0_usize, INFINITY),
  26.                             |(min_ix, min_dist), (ix, ci)| {
  27.                                 let dist = distance(xi, ci);
  28.                                 if dist < min_dist {
  29.                                     (ix, dist)
  30.                                 } else {
  31.                                     (min_ix, min_dist)
  32.                                 }
  33.                             },
  34.                         );
  35.                         argmin
  36.                     })
  37.                     .collect()
  38.             }
  40.             /// Recompute the centroids given the current clustering
  41.             fn recompute_centroids(
  42.                 xs: &[Vec<$kind>],
  43.                 clustering: &[usize],
  44.                 k: usize,
  45.             ) -> Vec<Vec<$kind>> {
  46.                 let ndims = xs[0].len();
  48.                 // NOTE: Kind of inefficient because we sweep all the data from each of the
  49.                 // k centroids.
  50.                 (0..k)
  51.                     .map(|cluster_ix| {
  52.                         let mut centroid: Vec<$kind> = vec![0.0; ndims];
  53.                         let mut n_cluster: $kind = 0.0;
  54.                         xs.iter().zip(clustering.iter()).for_each(|(xi, &zi)| {
  55.                             if zi == cluster_ix {
  56.                                 n_cluster += 1.0;
  57.                                 xi.iter().enumerate().for_each(|(j, &x_ij)| {
  58.                                     centroid[j] += x_ij;
  59.                                 });
  60.                             }
  61.                         });
  62.                         centroid.iter().map(|&c_j| c_j / n_cluster).collect()
  63.                     })
  64.                     .collect()
  65.             }
  67.             /// Assign the N D-dimensional data, `xs`, to `k` clusters using K-Means clustering
  68.             pub fn kmeans(xs: Vec<Vec<$kind>>, k: usize) -> Vec<usize> {
  69.                 assert!(xs.len() >= k);
  71.                 // Rather than pulling in a dependency to radomly select the staring
  72.                 // points for the centroids, we're going to deterministally choose them by
  73.                 // slecting evenly spaced points in `xs`
  74.                 let n_per_cluster: usize = xs.len() / k;
  75.                 let centroids: Vec<Vec<$kind>> =
  76.                     (0..k).map(|j| xs[j * n_per_cluster].clone()).collect();
  78.                 let mut clustering = nearest_centroids(&xs, &centroids);
  80.                 loop {
  81.                     let centroids = recompute_centroids(&xs, &clustering, k);
  82.                     let new_clustering = nearest_centroids(&xs, &centroids);
  84.                     // loop until the clustering doesn't change after the new centroids are computed
  85.                     if new_clustering
  86.                         .iter()
  87.                         .zip(clustering.iter())
  88.                         .all(|(&za, &zb)| za == zb)
  89.                     {
  90.                         // We need to use `return` to break out of the `loop`
  91.                         return clustering;
  92.                     } else {
  93.                         clustering = new_clustering;
  94.                     }
  95.                 }
  96.             }
  97.         }
  98.     };
  99. }
  101. // generate code for kmeans for f32 and f64 data
  102. impl_kmeans!(f64, f64);
  103. impl_kmeans!(f32, f32);
  105. #[cfg(test)]
  106. mod test {
  107.     use self::super::f64::kmeans;
  109.     #[test]
  110.     fn easy_univariate_clustering() {
  111.         let xs: Vec<Vec<f64>> = vec![
  112.             vec![-1.1],
  113.             vec![-1.2],
  114.             vec![-1.3],
  115.             vec![-1.4],
  116.             vec![1.1],
  117.             vec![1.2],
  118.             vec![1.3],
  119.             vec![1.4],
  120.         ];
  121.         let clustering = kmeans(xs, 2);
  122.         assert_eq!(clustering, vec![0, 0, 0, 0, 1, 1, 1, 1]);
  123.     }
  125.     #[test]
  126.     fn easy_univariate_clustering_odd_number_of_data() {
  127.         let xs: Vec<Vec<f64>> = vec![
  128.             vec![-1.1],
  129.             vec![-1.2],
  130.             vec![-1.3],
  131.             vec![-1.4],
  132.             vec![1.1],
  133.             vec![1.2],
  134.             vec![1.3],
  135.             vec![1.4],
  136.             vec![1.5],
  137.         ];
  138.         let clustering = kmeans(xs, 2);
  139.         assert_eq!(clustering, vec![0, 0, 0, 0, 1, 1, 1, 1, 1]);
  140.     }
  142.     #[test]
  143.     fn easy_bivariate_clustering() {
  144.         let xs: Vec<Vec<f64>> = vec![
  145.             vec![-1.1, 0.2],
  146.             vec![-1.2, 0.3],
  147.             vec![-1.3, 0.1],
  148.             vec![-1.4, 0.4],
  149.             vec![1.1, -1.1],
  150.             vec![1.2, -1.0],
  151.             vec![1.3, -1.2],
  152.             vec![1.4, -1.3],
  153.         ];
  154.         let clustering = kmeans(xs, 2);
  155.         assert_eq!(clustering, vec![0, 0, 0, 0, 1, 1, 1, 1]);
  156.     }
  158.     #[test]
  159.     fn high_dims() {
  160.         let xs: Vec<Vec<f64>> = vec![
  161.             vec![-2.7825343, -1.7604825, -5.5550113, -2.9752946, -2.7874138],
  162.             vec![-2.9847919, -3.8209332, -2.1531757, -2.2710119, -2.3582877],
  163.             vec![-3.0109320, -2.2366132, -2.8048492, -1.2632331, -4.5755581],
  164.             vec![-2.8432186, -1.0383805, -2.2022826, -2.7435962, -2.0013399],
  165.             vec![-2.6638082, -3.5520086, -1.3684702, -2.1562444, -1.3186447],
  166.             vec![1.7409171, 1.9687576, 4.7162628, 4.5743537, 3.7905611],
  167.             vec![3.2932369, 2.8508700, 2.5580937, 2.0437325, 4.2192562],
  168.             vec![2.5843321, 2.8329818, 2.1329531, 3.2562319, 2.4878733],
  169.             vec![2.1859638, 3.2880048, 3.7018615, 2.3641232, 1.6281994],
  170.             vec![2.6201773, 0.9006588, 2.6774097, 1.8188620, 1.6076493],
  171.         ];
  173.         let clustering = kmeans(xs, 2);
  174.         assert_eq!(clustering, vec![0, 0, 0, 0, 0, 1, 1, 1, 1, 1]);
  175.     }
  176. }