haar_lib/tree/
mod.rs

1//! 木に関するもの
2
3pub mod auxiliary_tree;
4pub mod centroid;
5pub mod centroid_decomposition;
6pub mod depth_query;
7pub mod euler_tour;
8pub mod hld;
9pub mod lca;
10pub mod rerooting;
11pub mod rooted_isomorphism;
12pub mod rooting;
13pub mod tree_dp;
14pub mod utils;
15
16/// [`Tree`]にもたせる辺の満たすトレイト。
17pub trait TreeEdgeTrait {
18    /// 辺の重みの型
19    type Weight;
20    /// 辺の始点を返す。
21    fn from(&self) -> usize;
22    /// 辺の終点を返す。
23    fn to(&self) -> usize;
24    /// 辺の重みを返す。
25    fn weight(&self) -> Self::Weight;
26    /// 逆辺を返す。
27    fn rev(self) -> Self;
28}
29
30/// 始点、終点、重み、番号をもつ木の辺
31#[derive(Clone, Debug)]
32pub struct TreeEdge<T, I> {
33    /// 始点
34    pub from: usize,
35    /// 終点
36    pub to: usize,
37    /// 重み
38    pub weight: T,
39    /// 辺の番号
40    pub index: I,
41}
42
43impl<T, I> TreeEdge<T, I> {
44    /// `from`から`to`への重さ`weight`、辺番号`index`をもつ有向辺を作る。
45    pub fn new(from: usize, to: usize, weight: T, index: I) -> Self {
46        Self {
47            from,
48            to,
49            weight,
50            index,
51        }
52    }
53}
54
55impl<T: Clone, I> TreeEdgeTrait for TreeEdge<T, I> {
56    type Weight = T;
57    #[inline]
58    fn from(&self) -> usize {
59        self.from
60    }
61    #[inline]
62    fn to(&self) -> usize {
63        self.to
64    }
65    #[inline]
66    fn weight(&self) -> Self::Weight {
67        self.weight.clone()
68    }
69    fn rev(mut self) -> Self {
70        std::mem::swap(&mut self.from, &mut self.to);
71        self
72    }
73}
74
75/// 木のノード
76#[derive(Clone, Debug, Default)]
77pub struct TreeNode<E> {
78    /// 親ノードへの辺
79    pub parent: Option<E>,
80    /// 子ノードへの辺
81    pub children: Vec<E>,
82}
83
84impl<E: TreeEdgeTrait> TreeNode<E> {
85    /// 隣接辺を列挙するイテレータを返す。
86    pub fn neighbors(&self) -> impl DoubleEndedIterator<Item = &E> {
87        self.children.iter().chain(self.parent.iter())
88    }
89
90    /// 隣接辺の個数を返す。
91    pub fn neighbors_size(&self) -> usize {
92        self.children.len() + self.parent.as_ref().map_or(0, |_| 1)
93    }
94}
95
96/// 非根付き木を構築する
97pub struct TreeBuilder<E> {
98    nodes: Vec<TreeNode<E>>,
99}
100
101impl<E: TreeEdgeTrait + Clone> TreeBuilder<E> {
102    /// 頂点数`size`の[`TreeBuilder`]を生成する。
103    pub fn new(size: usize) -> Self {
104        Self {
105            nodes: vec![
106                TreeNode {
107                    parent: None,
108                    children: vec![],
109                };
110                size
111            ],
112        }
113    }
114
115    /// [`Tree`]を作る。
116    pub fn build(self) -> Tree<E> {
117        Tree {
118            nodes: self.nodes,
119            root: None,
120        }
121    }
122}
123
124impl<E: TreeEdgeTrait + Clone> Extend<E> for TreeBuilder<E> {
125    fn extend<T: IntoIterator<Item = E>>(&mut self, iter: T) {
126        for e in iter {
127            self.nodes[e.from()].children.push(e.clone());
128            self.nodes[e.to()].children.push(e.rev());
129        }
130    }
131}
132
133/// 根付き木を構築する
134pub struct RootedTreeBuilder<E> {
135    nodes: Vec<TreeNode<E>>,
136    root: usize,
137}
138
139impl<E: TreeEdgeTrait + Clone> RootedTreeBuilder<E> {
140    /// 頂点数`size`の[`TreeBuilder`]を生成する。
141    pub fn new(size: usize, root: usize) -> Self {
142        Self {
143            nodes: vec![
144                TreeNode {
145                    parent: None,
146                    children: vec![],
147                };
148                size
149            ],
150            root,
151        }
152    }
153
154    /// 根付きの[`Tree`]を作る。
155    pub fn build(self) -> Tree<E> {
156        Tree {
157            nodes: self.nodes,
158            root: Some(self.root),
159        }
160    }
161}
162
163impl<E: TreeEdgeTrait + Clone> Extend<E> for RootedTreeBuilder<E> {
164    fn extend<T: IntoIterator<Item = E>>(&mut self, iter: T) {
165        for e in iter {
166            assert!(self.nodes[e.to()].parent.is_none());
167            self.nodes[e.from()].children.push(e.clone());
168            self.nodes[e.to()].parent.replace(e.rev());
169        }
170    }
171}
172
173/// 木
174#[derive(Clone, Debug)]
175pub struct Tree<E> {
176    nodes: Vec<TreeNode<E>>,
177    root: Option<usize>,
178}
179
180impl<E> Tree<E> {
181    /// 各頂点の[`TreeNode`]への参照のイテレータを返す。
182    pub fn nodes_iter(&self) -> impl Iterator<Item = &TreeNode<E>> {
183        self.nodes.iter()
184    }
185
186    /// `i`番目の頂点の[`TreeNode`]への参照を返す。
187    pub fn node_of(&self, i: usize) -> &TreeNode<E> {
188        &self.nodes[i]
189    }
190
191    /// 木の頂点数を返す。
192    pub fn len(&self) -> usize {
193        self.nodes.len()
194    }
195
196    /// 木の頂点数が`0`ならば`true`を返す。
197    pub fn is_empty(&self) -> bool {
198        self.nodes.is_empty()
199    }
200
201    /// 木に根があれば根を返す。
202    pub fn root(&self) -> Option<usize> {
203        self.root
204    }
205}