competitive_library/graph/
floyd_warshall.rs

1//! ワーシャルフロイト
2pub fn floyd_warshall(matrix: &[Vec<Option<i64>>]) -> Vec<Vec<Option<i64>>> {
3    let mut m: Vec<_> = matrix.to_vec();
4    let n = m.len();
5    (0..n).for_each(|i| {
6        (0..n).for_each(|j| {
7            (0..n).for_each(|k| {
8                m[j][k] = if m[j][k].is_none() && m[j][i].is_none() && m[i][k].is_none() {
9                    None
10                } else if m[j][i].is_none() || m[i][k].is_none() {
11                    m[j][k]
12                } else if m[j][k].is_none() {
13                    Some(m[j][i].unwrap() + m[i][k].unwrap())
14                } else {
15                    Some(std::cmp::min(
16                        m[j][k].unwrap(),
17                        m[j][i].unwrap() + m[i][k].unwrap(),
18                    ))
19                };
20            })
21        })
22    });
23    m
24}
25#[cfg(test)]
26mod tests {
27    use super::*;
28    #[test]
29    fn test_floyd_warshall_1() {
30        let matrix = vec![
31            vec![Some(0), Some(1), Some(1), None, Some(1)],
32            vec![Some(1), Some(0), Some(1), Some(1), None],
33            vec![Some(1), Some(1), Some(0), Some(1), Some(1)],
34            vec![None, Some(1), Some(1), Some(0), Some(1)],
35            vec![Some(1), None, Some(1), Some(1), Some(0)],
36        ];
37        let a = floyd_warshall(&matrix);
38
39        let ans = vec![
40            vec![Some(0), Some(1), Some(1), Some(2), Some(1)],
41            vec![Some(1), Some(0), Some(1), Some(1), Some(2)],
42            vec![Some(1), Some(1), Some(0), Some(1), Some(1)],
43            vec![Some(2), Some(1), Some(1), Some(0), Some(1)],
44            vec![Some(1), Some(2), Some(1), Some(1), Some(0)],
45        ];
46
47        assert_eq!(ans, a);
48    }
49
50    #[test]
51    fn test_floyd_warshall_2() {
52        let matrix = vec![
53            vec![Some(0), Some(10), None, Some(100)],
54            vec![None, Some(0), None, Some(1000)],
55            vec![None, Some(1), Some(0), Some(10000)],
56            vec![Some(5), None, None, Some(0)],
57        ];
58
59        let a = floyd_warshall(&matrix);
60        let ans = vec![
61            vec![Some(0), Some(10), None, Some(100)],
62            vec![Some(1005), Some(0), None, Some(1000)],
63            vec![Some(1006), Some(1), Some(0), Some(1001)],
64            vec![Some(5), Some(15), None, Some(0)],
65        ];
66        assert_eq!(ans, a);
67    }
68}
competitive_library/graph/floyd_warshall.rs

competitive_library/graph/
floyd_warshall.rs
