haar_lib/graph/
functional_graph.rs

1//! Functional Graph
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc357/tasks/abc357_e>
5
6use crate::ds::unionfind::UnionFind;
7
8/// [`FunctionalGraph`]を構築するための構造体。
9pub struct FunctionalGraphBuilder {
10    next: Vec<Option<usize>>,
11}
12
13/// 頂点の種類
14#[derive(Clone, Copy, PartialEq, Debug)]
15pub enum Kind {
16    /// 閉路を構成している頂点。
17    Loop,
18    /// `Loop`でも`Leaf`でもない頂点。
19    Branch,
20    /// 入次数が`0`の頂点。
21    Leaf,
22}
23
24impl Kind {
25    /// `Loop`ならば`true`を返す。
26    pub fn is_loop(self) -> bool {
27        matches!(self, Self::Loop)
28    }
29
30    /// `Branch`ならば`true`を返す。
31    pub fn is_branch(self) -> bool {
32        matches!(self, Self::Branch)
33    }
34
35    /// `Leaf`ならば`true`を返す。
36    pub fn is_leaf(self) -> bool {
37        matches!(self, Self::Leaf)
38    }
39}
40
41/// 連結成分
42#[derive(Clone, Default, Debug)]
43pub struct Group {
44    /// 閉路
45    pub loops: Vec<usize>,
46    /// 枝
47    pub branches: Vec<usize>,
48    /// 葉
49    pub leaves: Vec<usize>,
50}
51
52/// 各頂点の出次数が`1`である、`n`頂点`n`辺の有向グラフ
53pub struct FunctionalGraph {
54    next: Vec<usize>,
55    v_kind: Vec<Kind>,
56    group_index: Vec<usize>,
57    groups: Vec<Group>,
58    children: Vec<Vec<usize>>,
59}
60
61impl FunctionalGraphBuilder {
62    /// `n`頂点の空なグラフを用意する。
63    pub fn new(n: usize) -> Self {
64        Self {
65            next: vec![None; n],
66        }
67    }
68
69    /// `from`から`to`への有向辺を追加する。
70    pub fn add(&mut self, from: usize, to: usize) {
71        assert!(self.next[from].is_none());
72        self.next[from] = Some(to);
73    }
74
75    /// [`FunctionalGraph`]を構築する。
76    pub fn build(self) -> FunctionalGraph {
77        assert!(self.next.iter().all(|a| a.is_some()));
78
79        let next = self.next.into_iter().flatten().collect::<Vec<_>>();
80        let n = next.len();
81
82        let mut uf = UnionFind::new(n);
83        for (cur, &next) in next.iter().enumerate() {
84            uf.merge(cur, next);
85        }
86
87        let mut index = vec![0; n];
88        let g_num = index
89            .iter_mut()
90            .enumerate()
91            .filter_map(|(i, index)| (uf.root_of(i) == i).then_some(index))
92            .enumerate()
93            .map(|(i, index)| *index = i)
94            .count();
95
96        let mut groups = vec![Group::default(); g_num];
97        let mut group_index = vec![0; n];
98
99        let mut in_deg = vec![0; n];
100        for &next in &next {
101            in_deg[next] += 1;
102        }
103
104        let mut v_kind = vec![None; n];
105        let mut stack = in_deg
106            .iter()
107            .enumerate()
108            .filter_map(|(i, &d)| (d == 0).then_some(i))
109            .inspect(|&i| {
110                v_kind[i] = Some(Kind::Leaf);
111                group_index[i] = index[uf.root_of(i)];
112                groups[group_index[i]].leaves.push(i);
113            })
114            .collect::<Vec<_>>();
115
116        while let Some(cur) = stack.pop() {
117            if v_kind[cur].is_none() {
118                v_kind[cur] = Some(Kind::Branch);
119                group_index[cur] = index[uf.root_of(cur)];
120                groups[group_index[cur]].branches.push(cur);
121            }
122
123            let next = next[cur];
124            in_deg[next] -= 1;
125            if in_deg[next] == 0 {
126                stack.push(next);
127            }
128        }
129
130        for i in 0..n {
131            if in_deg[i] != 0 {
132                v_kind[i] = Some(Kind::Loop);
133                group_index[i] = index[uf.root_of(i)];
134                groups[group_index[i]].loops.push(i);
135            }
136        }
137
138        let mut children = vec![vec![]; n];
139        for (i, &p) in next.iter().enumerate() {
140            if !v_kind[i].unwrap().is_loop() {
141                children[p].push(i);
142            }
143        }
144
145        FunctionalGraph {
146            next,
147            v_kind: v_kind.into_iter().flatten().collect(),
148            group_index,
149            groups,
150            children,
151        }
152    }
153}
154
155impl FunctionalGraph {
156    /// 頂点`i`から辺を辿った次の頂点を返す。
157    pub fn next_of(&self, i: usize) -> usize {
158        self.next[i]
159    }
160
161    /// 頂点`i`の種類を返す。
162    pub fn kind_of(&self, i: usize) -> Kind {
163        self.v_kind[i]
164    }
165
166    /// 頂点`i`が属する連結成分に割り当てられた番号を返す。
167    pub fn group_index_of(&self, i: usize) -> usize {
168        self.group_index[i]
169    }
170
171    /// 頂点`i`の属する連結成分を返す。
172    pub fn group_of(&self, i: usize) -> &Group {
173        &self.groups[self.group_index_of(i)]
174    }
175
176    /// すべての連結成分への参照を返す。
177    pub fn groups(&self) -> &[Group] {
178        &self.groups
179    }
180
181    /// 閉路を切断して、根付き森として見たときの、頂点`i`の子頂点列を返す。
182    pub fn children(&self, i: usize) -> impl Iterator<Item = usize> + '_ {
183        self.children[i].iter().copied()
184    }
185}
haar_lib/graph/functional_graph.rs

haar_lib/graph/
functional_graph.rs
