1pub 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
16pub trait TreeEdgeTrait {
18 type Weight;
20 fn from(&self) -> usize;
22 fn to(&self) -> usize;
24 fn weight(&self) -> Self::Weight;
26 fn rev(self) -> Self;
28}
29
30#[derive(Clone, Debug)]
32pub struct TreeEdge<T, I> {
33 pub from: usize,
35 pub to: usize,
37 pub weight: T,
39 pub index: I,
41}
42
43impl<T, I> TreeEdge<T, I> {
44 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#[derive(Clone, Debug, Default)]
77pub struct TreeNode<E> {
78 pub parent: Option<E>,
80 pub children: Vec<E>,
82}
83
84impl<E: TreeEdgeTrait> TreeNode<E> {
85 pub fn neighbors(&self) -> impl DoubleEndedIterator<Item = &E> {
87 self.children.iter().chain(self.parent.iter())
88 }
89
90 pub fn neighbors_size(&self) -> usize {
92 self.children.len() + self.parent.as_ref().map_or(0, |_| 1)
93 }
94}
95
96pub struct TreeBuilder<E> {
98 nodes: Vec<TreeNode<E>>,
99}
100
101impl<E: TreeEdgeTrait + Clone> TreeBuilder<E> {
102 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 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
133pub struct RootedTreeBuilder<E> {
135 nodes: Vec<TreeNode<E>>,
136 root: usize,
137}
138
139impl<E: TreeEdgeTrait + Clone> RootedTreeBuilder<E> {
140 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 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#[derive(Clone, Debug)]
175pub struct Tree<E> {
176 nodes: Vec<TreeNode<E>>,
177 root: Option<usize>,
178}
179
180impl<E> Tree<E> {
181 pub fn nodes_iter(&self) -> impl Iterator<Item = &TreeNode<E>> {
183 self.nodes.iter()
184 }
185
186 pub fn node_of(&self, i: usize) -> &TreeNode<E> {
188 &self.nodes[i]
189 }
190
191 pub fn len(&self) -> usize {
193 self.nodes.len()
194 }
195
196 pub fn is_empty(&self) -> bool {
198 self.nodes.is_empty()
199 }
200
201 pub fn root(&self) -> Option<usize> {
203 self.root
204 }
205}