1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159
use ndarray::Array2; use num_traits::{Bounded, Signed, Zero}; use fixedbitset::FixedBitSet; use std::iter::Sum; /// Compute the maximum matching between two disjoints sets of vertices /// using the /// [Kuhn-Munkres algorithm](https://en.wikipedia.org/wiki/Hungarian_algorithm) /// (also known as Hungarian algorithm). /// /// The weights between the first and second sets are given into the /// `weights` square matrix. The first axis is indexed by the X set, /// and the second axis by the Y set. The return value is a pair with /// the total assignments weight, and a vector containing indices in /// the Y set for every vertex in the X set. /// /// This algorithm executes in O(n³) where n is the cardinality of the sets. /// /// # Panics /// /// This function will panic if the `weights` matrix is not a square matrix. pub fn kuhn_munkres<C>(weights: &Array2<C>) -> (C, Vec<usize>) where C: Bounded + Sum<C> + Zero + Signed + Ord + Copy, { // We call x the rows and y the columns. n is the size of the matrix. let n = weights.shape()[0]; assert_eq!(n, weights.shape()[1]); // xy represents matchings for x, yz matchings for y let mut xy: Vec<Option<usize>> = vec![None; n]; let mut yx: Vec<Option<usize>> = vec![None; n]; // lx is the labelling for x nodes, ly the labelling for y nodes. We start // with an acceptable labelling with the maximum possible values for lx // and 0 for ly. let mut lx: Vec<C> = weights .outer_iter() .map(|row| row.into_iter().max().unwrap()) .cloned() .collect::<Vec<_>>(); let mut ly: Vec<C> = vec![Zero::zero(); n]; // s, augmenting, and slack will be reset every time they are reused. augmenting // contains Some(prev) when the corresponding node belongs to the augmenting path. let mut s = FixedBitSet::with_capacity(n); let mut alternating = Vec::with_capacity(n); let mut slack = vec![Zero::zero(); n]; let mut slackx = Vec::with_capacity(n); for root in 0..n { alternating.clear(); alternating.resize(n, None); // Find y such that the path is augmented. This will be set when breaking for the // loop below. Above the loop is some code to initialize the search. let mut y = { s.clear(); s.insert(root); // Slack for a vertex y is, initially, the margin between the // sum of the labels of root and y, and the weight between root and y. // As we add x nodes to the alternating path, we update the slack to // represent the smallest margin between one of the x nodes and y. for y in 0..n { slack[y] = lx[root] + ly[y] - weights[[root, y]]; } slackx.clear(); slackx.resize(n, root); Some(loop { let mut delta = Bounded::max_value(); let mut x = 0; let mut y = 0; // Select one of the smallest slack delta and its edge (x, y) // for y not in the alternating path already. for yy in 0..n { if alternating[yy].is_none() && slack[yy] < delta { delta = slack[yy]; x = slackx[yy]; y = yy; } } debug_assert!(s.contains(x)); // If some slack has been found, remove it from x nodes in the // alternating path, and add it to y nodes in the alternating path. // The slack of y nodes outside the alternating path will be reduced // by this minimal slack as well. if delta > Zero::zero() { for x in s.ones() { lx[x] = lx[x] - delta; } for y in 0..n { if alternating[y].is_some() { ly[y] = ly[y] + delta; } else { slack[y] = slack[y] - delta; } } } debug_assert!(lx[x] + ly[y] == weights[[x, y]]); // Add (x, y) to the alternating path. alternating[y] = Some(x); if yx[y].is_none() { // We have found an augmenting path. break y; } // This y node had a predecessor, add it to the set of x nodes // in the augmenting path. let x = yx[y].unwrap(); debug_assert!(!s.contains(x)); s.insert(x); // Update slack because of the added vertex in s might contain a // greater slack than with previously inserted x nodes in the augmenting // path. for y in 0..n { if alternating[y].is_none() { let alternate_slack = lx[x] + ly[y] - weights[[x, y]]; if slack[y] > alternate_slack { slack[y] = alternate_slack; slackx[y] = x; } } } }) }; // Inverse edges along the augmenting path. while y.is_some() { let x = alternating[y.unwrap()].unwrap(); let prec = xy[x]; yx[y.unwrap()] = Some(x); xy[x] = y; y = prec; } } ( lx.into_iter().sum::<C>() + ly.into_iter().sum(), xy.into_iter().map(|v| v.unwrap()).collect::<Vec<_>>(), ) } /// Compute the minimum matching between two disjoints sets of vertices /// using the /// [Kuhn-Munkres algorithm](https://en.wikipedia.org/wiki/Hungarian_algorithm) /// (also known as Hungarian algorithm). /// /// The weights between the first and second sets are given into the /// `weights` square matrix. The first axis is indexed by the X set, /// and the second axis by the Y set. The return value is a pair with /// the total assignments weight, and a vector containing indices in /// the Y set for every vertex in the X set. /// /// This algorithm executes in O(n³) where n is the cardinality of the sets. /// /// # Panics /// /// This function will panic if the `weights` matrix is not a square matrix. pub fn kuhn_munkres_min<C>(weights: &Array2<C>) -> (C, Vec<usize>) where C: Bounded + Sum<C> + Zero + Signed + Ord + Copy, { let (total, assignments) = kuhn_munkres(&-weights.clone()); (-total, assignments) }