1#[macro_export]
6macro_rules! tree_mo_algorithm {
7 (vertices; $tree:expr, $root:expr, $paths:expr;
8 |$u:tt| $add:expr, |$v:tt| $remove:expr, |$i:tt| $answer:expr $(,)?) => {{
9 let tree = $tree;
10 let root: usize = $root;
11 let mut trace = Vec::new();
12 let tour = tree
13 .path_euler_tour_builder(root)
14 .build_with_trace(|u| trace.push(u));
15 let lca = tree.lca(root);
16 let (ranges, extras): (Vec<_>, Vec<_>) = ($paths)
17 .into_iter()
18 .map(|(mut u, mut v): (usize, usize)| {
19 if tour.vidx[u][0] > tour.vidx[v][0] {
20 ::std::mem::swap(&mut u, &mut v);
21 }
22 let p = lca.lca(u, v);
23 if p == u {
24 ((tour.vidx[u][0], tour.vidx[v][0] + 1), !0usize)
25 } else {
26 ((tour.vidx[u][1], tour.vidx[v][0] + 1), p)
27 }
28 })
29 .unzip();
30 $crate::tree_mo_algorithm!(@run &trace, tour.vidx.len(), &ranges;
31 |j| extras[j]; |$u| $add, |$v| $remove, |$i| $answer);
32 }};
33 (edges; $tree:expr, $root:expr, $paths:expr;
34 |$u:tt| $add:expr, |$v:tt| $remove:expr, |$i:tt| $answer:expr $(,)?) => {{
35 let tree = $tree;
36 let root: usize = $root;
37 let tour = tree.path_euler_tour_builder(root).build();
38 let lca = tree.lca(root);
39 let mut trace = vec![0; 2 * tour.eidx.len()];
40 for (e, &[first, after_last]) in tour.eidx.iter().enumerate() {
41 trace[first - 1] = e;
42 trace[after_last - 2] = e;
43 }
44 let ranges: Vec<_> = ($paths)
45 .into_iter()
46 .map(|(mut u, mut v): (usize, usize)| {
47 if tour.vidx[u][0] > tour.vidx[v][0] {
48 ::std::mem::swap(&mut u, &mut v);
49 }
50 if lca.lca(u, v) == u {
51 (tour.vidx[u][0], tour.vidx[v][0])
52 } else {
53 (tour.vidx[u][1] - 1, tour.vidx[v][0])
54 }
55 })
56 .collect();
57 $crate::tree_mo_algorithm!(@run &trace, tour.eidx.len(), &ranges;
58 |_| !0usize; |$u| $add, |$v| $remove, |$i| $answer);
59 }};
60 (@run $trace:expr, $elements:expr, $ranges:expr;
61 |$query:tt| $extra:expr; |$u:tt| $add:expr, |$v:tt| $remove:expr, |$i:tt| $answer:expr) => {{
62 let trace: &[usize] = $trace;
63 let mut used = vec![false; $elements];
64 $crate::mo_algorithm!(
65 $ranges,
66 (l, r),
67 |j| $crate::tree_mo_algorithm!(@toggle trace[j], used; |$u| $add, |$v| $remove),
68 |j| $crate::tree_mo_algorithm!(@toggle trace[j], used; |$u| $add, |$v| $remove),
69 |query_index| {
70 let $query: usize = query_index;
71 let p: usize = $extra;
72 if p != !0 {
73 let $u = p;
74 $add;
75 }
76 let $i: usize = query_index;
77 $answer;
78 if p != !0 {
79 let $v = p;
80 $remove;
81 }
82 },
83 );
84 }};
85 (@toggle $element:expr, $used:ident;
86 |$u:tt| $add:expr, |$v:tt| $remove:expr) => {{
87 let element = $element;
88 if $used[element] {
89 let $v = element;
90 $remove;
91 } else {
92 let $u = element;
93 $add;
94 }
95 $used[element] = !$used[element];
96 }};
97}
98
99#[cfg(test)]
100mod tests {
101 use crate::{
102 graph::{Graph, UndirectedSparseGraph},
103 tools::{Xorshift, testutil::exhaustive_sequences},
104 tree::{MixedTree, PathTree, StarTree},
105 };
106
107 fn verify(tree: &UndirectedSparseGraph, root: usize) {
108 let n = tree.vertices_size();
109 let mut expected = Vec::new();
110 for s in 0..n {
111 let mut parent = vec![n; n];
112 let mut edge = vec![n; n];
113 let mut queue = std::collections::VecDeque::from([s]);
114 parent[s] = s;
115 while let Some(u) = queue.pop_front() {
116 for a in tree.neighbors(u) {
117 if parent[a.to] == n {
118 parent[a.to] = u;
119 edge[a.to] = a.label;
120 queue.push_back(a.to);
121 }
122 }
123 }
124 for t in 0..n {
125 let mut vertices = vec![false; n];
126 let mut edges = vec![false; n - 1];
127 let mut u = t;
128 while u != s {
129 vertices[u] = true;
130 edges[edge[u]] = true;
131 u = parent[u];
132 }
133 vertices[s] = true;
134 expected.push((vertices, edges));
135 }
136 }
137 let mut count = 0;
138 tree_mo_algorithm!(vertices; tree, root, std::iter::once((root, root));
139 |_| count += 1, |_| count -= 1,
140 |_| assert_eq!(count, expected[root * n + root].0.iter().filter(|&&v| v).count()));
141 let mut paths: Vec<_> = (0..n).flat_map(|u| (0..n).map(move |v| (u, v))).collect();
142 paths.reverse();
143 paths.extend_from_within(..n);
144 for paths in [&[][..], &paths[..]] {
145 let mut present = vec![false; n];
146 let mut seen = vec![0; paths.len()];
147 tree_mo_algorithm!(vertices; tree, root, paths.iter().copied();
148 |u| { assert!(!present[u]); present[u] = true; },
149 |u| { assert!(present[u]); present[u] = false; },
150 |i| {
151 let (s, t) = paths[i];
152 assert_eq!(present, expected[s * n + t].0);
153 seen[i] += 1;
154 });
155 assert!(seen.iter().all(|&count| count == 1));
156 let mut present = vec![false; n - 1];
157 let mut seen = vec![0; paths.len()];
158 tree_mo_algorithm!(edges; tree, root, paths.iter().copied();
159 |e| { assert!(!present[e]); present[e] = true; },
160 |e| { assert!(present[e]); present[e] = false; },
161 |i| {
162 let (s, t) = paths[i];
163 assert_eq!(present, expected[s * n + t].1);
164 seen[i] += 1;
165 });
166 assert!(seen.iter().all(|&count| count == 1));
167 }
168 }
169
170 #[test]
171 fn test_tree_mo_algorithm() {
172 for n in 1..=6 {
173 for parents in exhaustive_sequences(0..n, n - 1..=n - 1) {
174 if !parents.iter().enumerate().all(|(i, _)| {
175 let mut u = i + 1;
176 for _ in 0..n {
177 if u == 0 {
178 return true;
179 }
180 u = parents[u - 1];
181 }
182 false
183 }) {
184 continue;
185 }
186 let tree = UndirectedSparseGraph::from_edges(
187 n,
188 parents
189 .into_iter()
190 .enumerate()
191 .map(|(i, p)| (p, i + 1))
192 .collect(),
193 );
194 for root in 0..n {
195 verify(&tree, root);
196 }
197 }
198 }
199 let mut rng = Xorshift::default();
200 for n in 1..=40 {
201 for tree in [
202 rng.random(PathTree(n)),
203 rng.random(StarTree(n)),
204 rng.random(MixedTree(n)),
205 ] {
206 verify(&tree, rng.random(0..n));
207 }
208 }
209 }
210}