munkres/
weight_matrix.rs

1use crate::{Position, SquareMatrix, WeightNum, Weights};
2
3#[derive(Debug)]
4pub struct WeightMatrix<T: WeightNum> {
5    c: SquareMatrix<T>,
6}
7
8impl<T: WeightNum> Weights for WeightMatrix<T> {
9    type T = T;
10
11    #[inline(always)]
12    fn n(&self) -> usize {
13        self.c.shape()[0]
14    }
15
16    #[inline]
17    fn element_at(&self, pos: Position) -> T {
18        self.c[(pos.row, pos.column)]
19    }
20
21    // for each row, subtracts the minimum of that row from each other value in the
22    // row.
23    fn sub_min_of_each_row(&mut self) {
24        for row in 0..self.n() {
25            let min = self.min_of_row(row);
26            self.sub_row(row, min);
27        }
28    }
29
30    // Add `val` to every element in row `row`.
31    fn add_row(&mut self, row: usize, val: T) {
32        self.c
33            .row_mut(row)
34            .mapv_inplace(|cur| cur.add_if_valid(val));
35    }
36
37    // Subtract `val` from every element in column `col`.
38    fn sub_column(&mut self, col: usize, val: T) {
39        self.c
40            .column_mut(col)
41            .mapv_inplace(|cur| cur.sub_if_valid(val));
42    }
43
44    fn is_solvable(&self) -> bool {
45        for row in self.c.genrows() {
46            if row.iter().all(|c| !c.is_valid()) {
47                return false;
48            }
49        }
50        true
51    }
52}
53
54impl<T: WeightNum> WeightMatrix<T> {
55    pub fn from_row_vec(n: usize, data: Vec<T>) -> WeightMatrix<T> {
56        WeightMatrix {
57            c: SquareMatrix::from_shape_vec((n, n), data).unwrap(),
58        }
59    }
60
61    pub fn from_fn<F: Fn((usize, usize)) -> T>(n: usize, f: F) -> WeightMatrix<T> {
62        assert!(n > 0);
63        WeightMatrix {
64            c: SquareMatrix::from_shape_fn((n, n), f),
65        }
66    }
67
68    /// Return the minimum element of row `row`.
69    fn min_of_row(&self, row: usize) -> T {
70        let row_iter = self.c.row(row);
71        let mut valid_iter = row_iter.iter().filter(|cost| cost.is_valid()).cloned();
72        let first_min = valid_iter.next().unwrap();
73        valid_iter.fold(
74            first_min,
75            |total_min, val| if val < total_min { val } else { total_min },
76        )
77    }
78
79    // Subtract `val` from every element in row `row`.
80    fn sub_row(&mut self, row: usize, val: T) {
81        self.c
82            .row_mut(row)
83            .mapv_inplace(|cur| cur.sub_if_valid(val));
84    }
85
86    pub fn as_slice(&self) -> &[T] {
87        self.c.as_slice().unwrap()
88    }
89}
90
91#[test]
92fn test_weight_matrix() {
93    assert_eq!(0, WeightMatrix::from_row_vec(1, vec![0]).min_of_row(0));
94    assert_eq!(1, WeightMatrix::from_row_vec(1, vec![1]).min_of_row(0));
95    assert_eq!(
96        1,
97        WeightMatrix::from_row_vec(2, vec![5, 1, 0, 0]).min_of_row(0)
98    );
99
100    let mut mat = WeightMatrix::from_row_vec(2, vec![0, 1, 2, 3]);
101    mat.sub_row(1, 1);
102    assert_eq!(&[0, 1, 1, 2], mat.as_slice());
103
104    let mut mat = WeightMatrix::from_row_vec(2, vec![5, 3, 2, 3]);
105    mat.sub_min_of_each_row();
106    assert_eq!(&[2, 0, 0, 1], mat.as_slice());
107}
munkres/weight_matrix.rs

munkres/
weight_matrix.rs
