dsalgo/
graph_bfs_path_count.rs

1pub fn path_count(
2    g: &[Vec<usize>],
3    src: usize,
4    modulus: usize,
5) -> Vec<usize> {
6    let n = g.len();
7
8    let mut f = vec![0; n];
9
10    f[src] = 1;
11
12    let mut que = std::collections::VecDeque::new();
13
14    que.push_back(src);
15
16    let mut to_que = vec![false; n];
17
18    to_que[src] = true;
19
20    while let Some(u) = que.pop_front() {
21        for &v in g[u].iter() {
22            f[v] += f[u];
23
24            f[v] %= modulus;
25
26            if to_que[v] {
27                continue;
28            }
29
30            to_que[v] = true;
31
32            que.push_back(v);
33        }
34    }
35
36    f
37}
38
39#[cfg(test)]
40
41mod tests {
42
43    use super::*;
44
45    #[test]
46
47    fn test_atcoder_abc244_e() {
48        let cases = vec![
49            ((4, 4, 4, 1, 3, 2, vec![(1, 2), (2, 3), (3, 4), (1, 4)]), 4),
50            (
51                (
52                    6,
53                    5,
54                    10,
55                    1,
56                    2,
57                    3,
58                    vec![(2, 3), (2, 4), (4, 6), (3, 6), (1, 5)],
59                ),
60                0,
61            ),
62            (
63                (
64                    10,
65                    15,
66                    20,
67                    4,
68                    4,
69                    6,
70                    vec![
71                        (2, 6),
72                        (2, 7),
73                        (5, 7),
74                        (4, 5),
75                        (2, 4),
76                        (3, 7),
77                        (1, 7),
78                        (1, 4),
79                        (2, 9),
80                        (5, 10),
81                        (1, 3),
82                        (7, 8),
83                        (7, 9),
84                        (1, 6),
85                        (1, 2),
86                    ],
87                ),
88                952504739,
89            ),
90        ];
91
92        for ((n, m, k, s, t, x, edges), ans) in cases {
93            let s = s - 1;
94
95            let t = t - 1;
96
97            let x = x - 1;
98
99            const B: usize = 1 << 11;
100
101            let encode =
102                |i: usize, u: usize, p: usize| -> usize { (p * n + u) * B + i };
103
104            const N: usize = 1 << 23;
105
106            let mut g = vec![vec![]; N];
107
108            for (u, v) in edges {
109                let u = u - 1;
110
111                let v = v - 1;
112
113                let p = if v == x { 1 } else { 0 };
114
115                let q = if u == x { 1 } else { 0 };
116
117                for i in 0..k {
118                    for j in 0..2 {
119                        let x = encode(i, u, j);
120
121                        let y = encode(i + 1, v, j ^ p);
122
123                        g[x].push(y);
124
125                        let x = encode(i, v, j);
126
127                        let y = encode(i + 1, u, j ^ q);
128
129                        g[x].push(y);
130                    }
131                }
132            }
133
134            let s = encode(0, s, 0);
135
136            let t = encode(k, t, 0);
137
138            const MOD: usize = 998_244_353;
139
140            let cnt = path_count(&g, s, MOD);
141
142            assert_eq!(cnt[t], ans);
143        }
144    }
145}
dsalgo/graph_bfs_path_count.rs

dsalgo/
graph_bfs_path_count.rs
