haar_lib/graph/
matrix_tree.rs

1//! 行列木定理を用いた数え上げ
2//! # Problems
3//! - <https://judge.yosupo.jp/problem/counting_spanning_tree_undirected>
4//! - <https://judge.yosupo.jp/problem/counting_spanning_tree_directed>
5
6use crate::graph::*;
7use crate::linalg::mod_p::determinant::*;
8use crate::math::prime_mod::PrimeMod;
9use crate::num::const_modint::*;
10
11/// 無向グラフにおいて、無向全域木の個数を数える。
12pub fn count_undirected_spanning_tree<P: PrimeMod>(
13    g: &Graph<Undirected, impl EdgeTrait>,
14) -> ConstModInt<P> {
15    let modulo = ConstModIntBuilder::<P>::new();
16
17    let n = g.len();
18    let mut lap = vec![vec![modulo.from_u64(0); n - 1]; n - 1];
19
20    for e in g.nodes_iter().flatten() {
21        let from = e.from();
22        let to = e.to();
23
24        if from < n - 1 {
25            lap[from][from] += modulo.from_u64(1);
26            if to < n - 1 {
27                lap[from][to] -= modulo.from_u64(1);
28            }
29        }
30    }
31
32    determinant(lap, &modulo)
33}
34
35/// 有向グラフにおいて、頂点`root`を終点根とするような、有向全域木の個数を数える。
36pub fn count_directed_spanning_tree<P: PrimeMod>(
37    g: &Graph<Directed, impl EdgeTrait>,
38    root: usize,
39) -> ConstModInt<P> {
40    let modulo = ConstModIntBuilder::<P>::new();
41
42    let n = g.len();
43    let mut lap = vec![vec![modulo.from_u64(0); n - 1]; n - 1];
44
45    for e in g.nodes_iter().flatten() {
46        let mut from = e.from();
47        let mut to = e.to();
48
49        if from == root {
50            continue;
51        }
52        if from > root {
53            from -= 1;
54        }
55        lap[from][from] += modulo.from_u64(1);
56
57        if to == root {
58            continue;
59        }
60        if to > root {
61            to -= 1;
62        }
63        lap[from][to] -= modulo.from_u64(1);
64    }
65
66    determinant(lap, &modulo)
67}
haar_lib/graph/matrix_tree.rs

haar_lib/graph/
matrix_tree.rs
