1use super::{Graph, Magma, Monoid, UndirectedSparseGraph, Unital};
2use std::mem::MaybeUninit;
3
4pub trait MonoidCluster {
5 type Vertex;
6 type Edge;
7 type PathMonoid: Monoid;
8 type PointMonoid: Monoid;
9
10 fn add_edge(
11 path: &<<Self as MonoidCluster>::PathMonoid as Magma>::T,
12 ) -> <<Self as MonoidCluster>::PointMonoid as Magma>::T;
13 fn add_vertex(
14 point: &<<Self as MonoidCluster>::PointMonoid as Magma>::T,
15 vertex: &Self::Vertex,
16 parent_edge: Option<&Self::Edge>,
17 ) -> <<Self as MonoidCluster>::PathMonoid as Magma>::T;
18}
19
20pub trait Cluster {
21 type Vertex;
22 type Edge;
23 type Path: Clone;
24 type Point: Clone;
25
26 fn unit_path() -> Self::Path;
27 fn unit_point() -> Self::Point;
28 fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path;
29 fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point;
30 fn add_edge(path: &Self::Path) -> Self::Point;
31 fn add_vertex(
32 point: &Self::Point,
33 vertex: &Self::Vertex,
34 parent_edge: Option<&Self::Edge>,
35 ) -> Self::Path;
36
37 fn vertex(vertex: &Self::Vertex, parent_edge: Option<&Self::Edge>) -> Self::Path {
38 Self::add_vertex(&Self::unit_point(), vertex, parent_edge)
39 }
40}
41
42impl<C> Cluster for C
43where
44 C: MonoidCluster,
45{
46 type Vertex = C::Vertex;
47 type Edge = C::Edge;
48 type Path = <<C as MonoidCluster>::PathMonoid as Magma>::T;
49 type Point = <<C as MonoidCluster>::PointMonoid as Magma>::T;
50
51 fn unit_path() -> Self::Path {
52 <C::PathMonoid as Unital>::unit()
53 }
54
55 fn unit_point() -> Self::Point {
56 <C::PointMonoid as Unital>::unit()
57 }
58
59 fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path {
60 <C::PathMonoid as Magma>::operate(left, right)
61 }
62
63 fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point {
64 <C::PointMonoid as Magma>::operate(left, right)
65 }
66
67 fn add_edge(path: &Self::Path) -> Self::Point {
68 <C as MonoidCluster>::add_edge(path)
69 }
70
71 fn add_vertex(
72 point: &Self::Point,
73 vertex: &Self::Vertex,
74 parent_edge: Option<&Self::Edge>,
75 ) -> Self::Path {
76 <C as MonoidCluster>::add_vertex(point, vertex, parent_edge)
77 }
78}
79
80#[derive(Clone)]
81pub struct StaticTopTree {
82 root: usize,
83 n: usize,
84 edge_child: Vec<usize>,
85 parent_edge: Vec<usize>,
86 compressed: Vec<InnerNode>,
87 raked: Vec<InnerNode>,
88 vertex_links: Vec<VertexLinks>,
89 compress_roots: Vec<Option<Slot>>,
90 rake_roots: Vec<Option<Slot>>,
91}
92
93#[derive(Clone)]
94struct InnerNode {
95 left: Slot,
96 right: Slot,
97 parent: usize,
98}
99
100#[derive(Clone)]
101struct InnerValue<T> {
102 parent: usize,
103 left: T,
104 right: T,
105}
106
107pub struct StaticTopTreeDp<'a, C>
108where
109 C: Cluster,
110{
111 tree: &'a StaticTopTree,
112 vertices: Vec<<C as Cluster>::Vertex>,
113 edges: Vec<<C as Cluster>::Edge>,
114 compressed: Vec<InnerValue<<C as Cluster>::Path>>,
115 raked: Vec<InnerValue<<C as Cluster>::Point>>,
116 light_points: Vec<<C as Cluster>::Point>,
117 all_point: <C as Cluster>::Point,
118}
119
120#[derive(Debug, Clone, Copy)]
121struct VertexLinks {
122 heavy_parent: usize,
123 compress_parent: usize,
124 rake_parent: usize,
125}
126
127#[derive(Debug)]
128struct Node {
129 depth: usize,
130 slot: Slot,
131}
132
133#[derive(Debug, Clone, Copy)]
134enum Slot {
135 CompressLeaf(usize),
136 CompressInner(usize),
137 RakeLeaf(usize),
138 RakeInner(usize),
139}
140
141struct RootedInfo {
142 order: Vec<usize>,
143 children_start: Vec<usize>,
144 children: Vec<usize>,
145 edge_child: Vec<usize>,
146 parent_edge: Vec<usize>,
147}
148
149impl UndirectedSparseGraph {
150 pub fn static_top_tree(&self, root: usize) -> StaticTopTree {
151 StaticTopTree::new(root, self)
152 }
153}
154
155impl StaticTopTree {
156 pub fn new(root: usize, graph: &UndirectedSparseGraph) -> Self {
157 let n = graph.vertices_size();
158 assert!(n > 0);
159 assert!(root < n);
160 assert_eq!(graph.edges_size() + 1, n);
161
162 let RootedInfo {
163 order,
164 children_start,
165 children,
166 edge_child,
167 parent_edge,
168 } = rooted_children(graph, root);
169 let mut this = Self {
170 root,
171 n,
172 edge_child,
173 parent_edge,
174 compressed: Vec::with_capacity(n.saturating_sub(1)),
175 raked: Vec::with_capacity(n.saturating_sub(1)),
176 vertex_links: vec![
177 VertexLinks {
178 heavy_parent: usize::MAX,
179 compress_parent: usize::MAX,
180 rake_parent: usize::MAX,
181 };
182 n
183 ],
184 compress_roots: vec![None; n],
185 rake_roots: vec![None; n],
186 };
187
188 let mut heavy_child = vec![usize::MAX; n];
189 let mut mask = vec![1u64; n];
190 let mut buckets: [Vec<Node>; 64] = std::array::from_fn(|_| Vec::new());
191
192 for &u in order.iter().rev() {
193 let children = &children[children_start[u]..children_start[u + 1]];
194 let mut sum_rake = 0u64;
195 for &v in children {
196 sum_rake += bit_ceil(mask[v]) << 1;
197 }
198 mask[u] = bit_ceil(sum_rake);
199 for &v in children {
200 let child = bit_ceil(mask[v]) << 1;
201 let depth = bit_ceil(sum_rake - child).trailing_zeros() as usize;
202 let step = 1u64 << depth;
203 let cand = ((mask[v] + step - 1) >> depth << depth) + step;
204 if cand <= mask[u] {
205 mask[u] = cand;
206 heavy_child[u] = v;
207 }
208 }
209
210 let mut has = 0u64;
211 let mut num_light = 0usize;
212 for &v in children {
213 if v == heavy_child[u] {
214 continue;
215 }
216 num_light += 1;
217 let child = bit_ceil(mask[v]) << 1;
218 let depth = bit_ceil(sum_rake - child).trailing_zeros() as usize;
219 this.build_compress(v, &heavy_child, &mask);
220 buckets[depth].push(Node {
221 depth,
222 slot: Slot::RakeLeaf(v),
223 });
224 has |= 1u64 << depth;
225 }
226 if num_light == 0 {
227 continue;
228 }
229
230 while num_light > 1 {
231 let left = pop_bucket(&mut buckets, &mut has);
232 let right = pop_bucket(&mut buckets, &mut has);
233 let node = this.merge_rake(left, right);
234 let depth = node.depth;
235 buckets[depth].push(node);
236 has |= 1u64 << depth;
237 num_light -= 1;
238 }
239
240 let root = pop_bucket(&mut buckets, &mut has);
241 this.rake_roots[u] = Some(root.slot);
242 for &v0 in children {
243 if v0 == heavy_child[u] {
244 continue;
245 }
246 let rake_parent = this.vertex_links[v0].rake_parent;
247 let mut v = v0;
248 while v != usize::MAX {
249 this.vertex_links[v].heavy_parent = u;
250 this.vertex_links[v].rake_parent = rake_parent;
251 v = heavy_child[v];
252 }
253 }
254 }
255
256 this.build_compress(root, &heavy_child, &mask);
257 this
258 }
259
260 pub fn vertices_size(&self) -> usize {
261 self.n
262 }
263
264 pub fn edges_size(&self) -> usize {
265 self.edge_child.len()
266 }
267
268 pub fn dp<C>(
269 &self,
270 vertices: Vec<<C as Cluster>::Vertex>,
271 edges: Vec<<C as Cluster>::Edge>,
272 ) -> StaticTopTreeDp<'_, C>
273 where
274 C: Cluster,
275 {
276 StaticTopTreeDp::new(self, vertices, edges)
277 }
278
279 pub fn fold_all<C>(
280 &self,
281 vertices: &[<C as Cluster>::Vertex],
282 edges: &[<C as Cluster>::Edge],
283 ) -> <C as Cluster>::Point
284 where
285 C: Cluster,
286 {
287 assert_eq!(vertices.len(), self.vertices_size());
288 assert_eq!(edges.len(), self.edges_size());
289 let path = self.fold_compress::<C>(
290 vertices,
291 edges,
292 self.compress_roots[self.root].expect("root compress tree must exist"),
293 );
294 C::add_edge(&path)
295 }
296
297 fn build_compress(&mut self, mut vertex: usize, heavy_child: &[usize], mask: &[u64]) -> Node {
298 let start = vertex;
299 let mut stack = Vec::new();
300 while vertex != usize::MAX {
301 stack.push(Node {
302 depth: bit_ceil(mask[vertex]).trailing_zeros() as usize,
303 slot: Slot::CompressLeaf(vertex),
304 });
305 loop {
306 let len = stack.len();
307 if len >= 3
308 && (stack[len - 3].depth == stack[len - 2].depth
309 || stack[len - 3].depth <= stack[len - 1].depth)
310 {
311 let tail = stack.pop().unwrap();
312 let right = stack.pop().unwrap();
313 let left = stack.pop().unwrap();
314 let node = self.merge_compress(left, right);
315 stack.push(node);
316 stack.push(tail);
317 } else if len >= 2 && stack[len - 2].depth <= stack[len - 1].depth {
318 let right = stack.pop().unwrap();
319 let left = stack.pop().unwrap();
320 stack.push(self.merge_compress(left, right));
321 } else {
322 break;
323 }
324 }
325 vertex = heavy_child[vertex];
326 }
327 while stack.len() > 1 {
328 let right = stack.pop().unwrap();
329 let left = stack.pop().unwrap();
330 stack.push(self.merge_compress(left, right));
331 }
332 let root = stack.pop().unwrap();
333 self.compress_roots[start] = Some(root.slot);
334 root
335 }
336
337 fn merge_compress(&mut self, left: Node, right: Node) -> Node {
338 let id = self.compressed.len();
339 self.set_parent(left.slot, id << 1);
340 self.set_parent(right.slot, id << 1 | 1);
341 self.compressed.push(InnerNode {
342 left: left.slot,
343 right: right.slot,
344 parent: usize::MAX,
345 });
346 Node {
347 depth: left.depth.max(right.depth) + 1,
348 slot: Slot::CompressInner(id),
349 }
350 }
351
352 fn merge_rake(&mut self, left: Node, right: Node) -> Node {
353 let id = self.raked.len();
354 self.set_parent(left.slot, id << 1);
355 self.set_parent(right.slot, id << 1 | 1);
356 self.raked.push(InnerNode {
357 left: left.slot,
358 right: right.slot,
359 parent: usize::MAX,
360 });
361 Node {
362 depth: left.depth.max(right.depth) + 1,
363 slot: Slot::RakeInner(id),
364 }
365 }
366
367 fn set_parent(&mut self, slot: Slot, parent: usize) {
368 match slot {
369 Slot::CompressLeaf(v) => self.vertex_links[v].compress_parent = parent,
370 Slot::CompressInner(i) => self.compressed[i].parent = parent,
371 Slot::RakeLeaf(v) => self.vertex_links[v].rake_parent = parent,
372 Slot::RakeInner(i) => self.raked[i].parent = parent,
373 }
374 }
375
376 fn init_compress<C>(
377 &self,
378 data: &mut StaticTopTreeDataBuilder<C>,
379 vertices: &[<C as Cluster>::Vertex],
380 edges: &[<C as Cluster>::Edge],
381 slot: Slot,
382 ) -> <C as Cluster>::Path
383 where
384 C: Cluster,
385 {
386 match slot {
387 Slot::CompressLeaf(vertex) => {
388 let point = self.init_point(data, vertices, edges, vertex);
389 C::add_vertex(
390 &point,
391 &vertices[vertex],
392 self.parent_edge_ref(edges, vertex),
393 )
394 }
395 Slot::CompressInner(id) => {
396 let node = &self.compressed[id];
397 let left = self.init_compress(data, vertices, edges, node.left);
398 let right = self.init_compress(data, vertices, edges, node.right);
399 data.compressed[id].write(InnerValue {
400 parent: node.parent,
401 left: left.clone(),
402 right: right.clone(),
403 });
404 C::compress(&left, &right)
405 }
406 Slot::RakeLeaf(_) | Slot::RakeInner(_) => unreachable!(),
407 }
408 }
409
410 fn fold_compress<C>(
411 &self,
412 vertices: &[<C as Cluster>::Vertex],
413 edges: &[<C as Cluster>::Edge],
414 slot: Slot,
415 ) -> <C as Cluster>::Path
416 where
417 C: Cluster,
418 {
419 match slot {
420 Slot::CompressLeaf(vertex) => {
421 let point = self.fold_point::<C>(vertices, edges, vertex);
422 C::add_vertex(
423 &point,
424 &vertices[vertex],
425 self.parent_edge_ref(edges, vertex),
426 )
427 }
428 Slot::CompressInner(id) => {
429 let node = &self.compressed[id];
430 let left = self.fold_compress::<C>(vertices, edges, node.left);
431 let right = self.fold_compress::<C>(vertices, edges, node.right);
432 C::compress(&left, &right)
433 }
434 Slot::RakeLeaf(_) | Slot::RakeInner(_) => unreachable!(),
435 }
436 }
437
438 fn init_point<C>(
439 &self,
440 data: &mut StaticTopTreeDataBuilder<C>,
441 vertices: &[<C as Cluster>::Vertex],
442 edges: &[<C as Cluster>::Edge],
443 vertex: usize,
444 ) -> <C as Cluster>::Point
445 where
446 C: Cluster,
447 {
448 let point = if let Some(slot) = self.rake_roots[vertex] {
449 self.init_rake(data, vertices, edges, slot)
450 } else {
451 C::unit_point()
452 };
453 data.light_points[vertex] = point.clone();
454 point
455 }
456
457 fn fold_point<C>(
458 &self,
459 vertices: &[<C as Cluster>::Vertex],
460 edges: &[<C as Cluster>::Edge],
461 vertex: usize,
462 ) -> <C as Cluster>::Point
463 where
464 C: Cluster,
465 {
466 if let Some(slot) = self.rake_roots[vertex] {
467 self.fold_rake::<C>(vertices, edges, slot)
468 } else {
469 C::unit_point()
470 }
471 }
472
473 fn init_rake<C>(
474 &self,
475 data: &mut StaticTopTreeDataBuilder<C>,
476 vertices: &[<C as Cluster>::Vertex],
477 edges: &[<C as Cluster>::Edge],
478 slot: Slot,
479 ) -> <C as Cluster>::Point
480 where
481 C: Cluster,
482 {
483 match slot {
484 Slot::RakeLeaf(vertex) => {
485 let path = self.init_compress(
486 data,
487 vertices,
488 edges,
489 self.compress_roots[vertex].expect("light child path must exist"),
490 );
491 C::add_edge(&path)
492 }
493 Slot::RakeInner(id) => {
494 let node = &self.raked[id];
495 let left = self.init_rake(data, vertices, edges, node.left);
496 let right = self.init_rake(data, vertices, edges, node.right);
497 data.raked[id].write(InnerValue {
498 parent: node.parent,
499 left: left.clone(),
500 right: right.clone(),
501 });
502 C::rake(&left, &right)
503 }
504 Slot::CompressLeaf(_) | Slot::CompressInner(_) => unreachable!(),
505 }
506 }
507
508 fn fold_rake<C>(
509 &self,
510 vertices: &[<C as Cluster>::Vertex],
511 edges: &[<C as Cluster>::Edge],
512 slot: Slot,
513 ) -> <C as Cluster>::Point
514 where
515 C: Cluster,
516 {
517 match slot {
518 Slot::RakeLeaf(vertex) => {
519 let path = self.fold_compress::<C>(
520 vertices,
521 edges,
522 self.compress_roots[vertex].expect("light child path must exist"),
523 );
524 C::add_edge(&path)
525 }
526 Slot::RakeInner(id) => {
527 let node = &self.raked[id];
528 let left = self.fold_rake::<C>(vertices, edges, node.left);
529 let right = self.fold_rake::<C>(vertices, edges, node.right);
530 C::rake(&left, &right)
531 }
532 Slot::CompressLeaf(_) | Slot::CompressInner(_) => unreachable!(),
533 }
534 }
535
536 fn parent_edge_ref<'a, T>(&self, edges: &'a [T], vertex: usize) -> Option<&'a T> {
537 let edge = self.parent_edge[vertex];
538 if edge == usize::MAX {
539 None
540 } else {
541 Some(&edges[edge])
542 }
543 }
544}
545
546impl<'a, C> StaticTopTreeDp<'a, C>
547where
548 C: Cluster,
549{
550 pub fn new(
551 tree: &'a StaticTopTree,
552 vertices: Vec<<C as Cluster>::Vertex>,
553 edges: Vec<<C as Cluster>::Edge>,
554 ) -> Self {
555 assert_eq!(vertices.len(), tree.vertices_size());
556 assert_eq!(edges.len(), tree.edges_size());
557
558 let mut data: StaticTopTreeDataBuilder<C> = StaticTopTreeDataBuilder::new(tree);
559 let path = tree.init_compress(
560 &mut data,
561 &vertices,
562 &edges,
563 tree.compress_roots[tree.root].expect("root compress tree must exist"),
564 );
565 let all_point = C::add_edge(&path);
566 Self {
567 tree,
568 vertices,
569 edges,
570 compressed: unsafe { assume_init_vec(data.compressed) },
571 raked: unsafe { assume_init_vec(data.raked) },
572 light_points: data.light_points,
573 all_point,
574 }
575 }
576
577 pub fn get_vertex(&self, vertex: usize) -> &<C as Cluster>::Vertex {
578 &self.vertices[vertex]
579 }
580
581 pub fn apply_vertex<F>(&mut self, vertex: usize, f: F)
582 where
583 F: FnOnce(&mut <C as Cluster>::Vertex),
584 {
585 assert!(vertex < self.vertices.len());
586 f(&mut self.vertices[vertex]);
587 self.update_from_vertex(vertex);
588 }
589
590 pub fn set_vertex(&mut self, vertex: usize, value: <C as Cluster>::Vertex) {
591 self.apply_vertex(vertex, |x| *x = value);
592 }
593
594 pub fn get_edge(&self, edge: usize) -> &<C as Cluster>::Edge {
595 &self.edges[edge]
596 }
597
598 pub fn apply_edge<F>(&mut self, edge: usize, f: F)
599 where
600 F: FnOnce(&mut <C as Cluster>::Edge),
601 {
602 assert!(edge < self.edges.len());
603 f(&mut self.edges[edge]);
604 self.update_from_vertex(self.tree.edge_child[edge]);
605 }
606
607 pub fn set_edge(&mut self, edge: usize, value: <C as Cluster>::Edge) {
608 self.apply_edge(edge, |x| *x = value);
609 }
610
611 pub fn fold_all(&self) -> &<C as Cluster>::Point {
612 &self.all_point
613 }
614
615 #[inline(always)]
616 pub fn fold_path(&self, mut vertex: usize) -> <C as Cluster>::Path {
617 assert!(vertex < self.tree.n);
618 let mut path = C::unit_path();
619 let mut point = self.light_points[vertex].clone();
620 loop {
621 let links = self.tree.vertex_links[vertex];
622 let mut left = C::unit_path();
623 let mut right = C::unit_path();
624 let mut compress_parent = links.compress_parent;
625 while compress_parent != usize::MAX {
626 let inner = &self.compressed[compress_parent / 2];
627 if compress_parent & 1 == 0 {
628 right = C::compress(&right, &inner.right);
629 } else {
630 left = C::compress(&inner.left, &left);
631 }
632 compress_parent = inner.parent;
633 }
634 let right_point = C::add_edge(&right);
635 point = C::rake(&point, &right_point);
636 let mid = C::add_vertex(
637 &point,
638 &self.vertices[vertex],
639 self.tree.parent_edge_ref(&self.edges, vertex),
640 );
641 let mid = C::compress(&mid, &path);
642 path = C::compress(&left, &mid);
643 if links.heavy_parent == usize::MAX {
644 return path;
645 }
646
647 point = C::unit_point();
648 let mut rake_parent = links.rake_parent;
649 while rake_parent != usize::MAX {
650 let inner = &self.raked[rake_parent / 2];
651 if rake_parent & 1 == 0 {
652 point = C::rake(&point, &inner.right);
653 } else {
654 point = C::rake(&inner.left, &point);
655 }
656 rake_parent = inner.parent;
657 }
658 vertex = links.heavy_parent;
659 }
660 }
661
662 fn update_from_vertex(&mut self, mut vertex: usize) {
663 assert!(vertex < self.tree.n);
664 while vertex != usize::MAX {
665 let links = self.tree.vertex_links[vertex];
666 let base = C::add_vertex(
667 &self.light_points[vertex],
668 &self.vertices[vertex],
669 self.tree.parent_edge_ref(&self.edges, vertex),
670 );
671 let path = self.update_compress(links.compress_parent, base);
672 let point = C::add_edge(&path);
673 let point = self.update_rake(links.rake_parent, point);
674 if links.heavy_parent == usize::MAX {
675 self.all_point = point;
676 } else {
677 self.light_points[links.heavy_parent] = point;
678 }
679 vertex = links.heavy_parent;
680 }
681 }
682
683 fn update_compress(
684 &mut self,
685 mut id: usize,
686 mut path: <C as Cluster>::Path,
687 ) -> <C as Cluster>::Path {
688 while id != usize::MAX {
689 let inner = &mut self.compressed[id / 2];
690 if id & 1 == 0 {
691 inner.left = path;
692 } else {
693 inner.right = path;
694 }
695 path = C::compress(&inner.left, &inner.right);
696 id = inner.parent;
697 }
698 path
699 }
700
701 fn update_rake(
702 &mut self,
703 mut id: usize,
704 mut point: <C as Cluster>::Point,
705 ) -> <C as Cluster>::Point {
706 while id != usize::MAX {
707 let inner = &mut self.raked[id / 2];
708 if id & 1 == 0 {
709 inner.left = point;
710 } else {
711 inner.right = point;
712 }
713 point = C::rake(&inner.left, &inner.right);
714 id = inner.parent;
715 }
716 point
717 }
718}
719
720struct StaticTopTreeDataBuilder<C>
721where
722 C: Cluster,
723{
724 compressed: Vec<MaybeUninit<InnerValue<<C as Cluster>::Path>>>,
725 raked: Vec<MaybeUninit<InnerValue<<C as Cluster>::Point>>>,
726 light_points: Vec<<C as Cluster>::Point>,
727}
728
729impl<C> StaticTopTreeDataBuilder<C>
730where
731 C: Cluster,
732{
733 fn new(tree: &StaticTopTree) -> Self {
734 let mut compressed = Vec::with_capacity(tree.compressed.len());
735 compressed.resize_with(tree.compressed.len(), MaybeUninit::uninit);
736 let mut raked = Vec::with_capacity(tree.raked.len());
737 raked.resize_with(tree.raked.len(), MaybeUninit::uninit);
738 Self {
739 compressed,
740 raked,
741 light_points: vec![C::unit_point(); tree.n],
742 }
743 }
744}
745
746unsafe fn assume_init_vec<T>(mut vec: Vec<MaybeUninit<T>>) -> Vec<T> {
747 let len = vec.len();
748 let cap = vec.capacity();
749 let ptr = vec.as_mut_ptr() as *mut T;
750 std::mem::forget(vec);
751 unsafe { Vec::from_raw_parts(ptr, len, cap) }
752}
753
754fn bit_ceil(x: u64) -> u64 {
755 if x <= 1 { 1 } else { x.next_power_of_two() }
756}
757
758fn rooted_children(graph: &UndirectedSparseGraph, root: usize) -> RootedInfo {
759 let n = graph.vertices_size();
760 let mut order = Vec::with_capacity(n);
761 let mut parent = vec![usize::MAX; n];
762 let mut parent_edge = vec![usize::MAX; n];
763 let mut edge_child = vec![0; graph.edges_size()];
764 order.push(root);
765 parent[root] = usize::MAX;
766 for i in 0..n {
767 let u = order[i];
768 for a in graph.neighbors(u) {
769 if a.to == parent[u] {
770 continue;
771 }
772 parent[a.to] = u;
773 parent_edge[a.to] = a.label;
774 edge_child[a.label] = a.to;
775 order.push(a.to);
776 }
777 }
778 let mut children_start = vec![0usize; n + 1];
779 for &v in order.iter().skip(1) {
780 children_start[parent[v] + 1] += 1;
781 }
782 for i in 1..=n {
783 children_start[i] += children_start[i - 1];
784 }
785 let mut children = vec![0; n.saturating_sub(1)];
786 let mut child_pos = children_start.clone();
787 for &v in order.iter().skip(1) {
788 let pos = child_pos[parent[v]];
789 children[pos] = v;
790 child_pos[parent[v]] += 1;
791 }
792 RootedInfo {
793 order,
794 children_start,
795 children,
796 edge_child,
797 parent_edge,
798 }
799}
800
801fn pop_bucket(buckets: &mut [Vec<Node>; 64], has: &mut u64) -> Node {
802 let depth = has.trailing_zeros() as usize;
803 let node = buckets[depth].pop().unwrap();
804 if buckets[depth].is_empty() {
805 *has &= !(1u64 << depth);
806 }
807 node
808}
809
810#[cfg(test)]
811mod tests {
812 use super::*;
813 use crate::{
814 algebra::{Associative, Magma, Unital},
815 graph::UndirectedSparseGraph,
816 num::{One, Zero, mint_basic::MInt998244353},
817 tools::Xorshift,
818 tree::{MixedTree, PathTree, StarTree},
819 };
820 use std::array;
821
822 type MInt = MInt998244353;
823
824 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
825 struct Point {
826 sum: MInt,
827 cnt: MInt,
828 }
829
830 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
831 struct Path {
832 a: MInt,
833 b: MInt,
834 sum: MInt,
835 cnt: MInt,
836 }
837
838 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
839 struct PathPair {
840 forward: Path,
841 reverse: Path,
842 }
843
844 struct PointMonoid;
845 impl Magma for PointMonoid {
846 type T = Point;
847 fn operate(x: &Self::T, y: &Self::T) -> Self::T {
848 Point {
849 sum: x.sum + y.sum,
850 cnt: x.cnt + y.cnt,
851 }
852 }
853 }
854 impl Unital for PointMonoid {
855 fn unit() -> Self::T {
856 Point {
857 sum: MInt::zero(),
858 cnt: MInt::zero(),
859 }
860 }
861 }
862 impl Associative for PointMonoid {}
863
864 struct PathMonoid;
865 impl Magma for PathMonoid {
866 type T = Path;
867 fn operate(x: &Self::T, y: &Self::T) -> Self::T {
868 Path {
869 a: x.a * y.a,
870 b: x.b + x.a * y.b,
871 sum: x.sum + x.a * y.sum + x.b * y.cnt,
872 cnt: x.cnt + y.cnt,
873 }
874 }
875 }
876 impl Unital for PathMonoid {
877 fn unit() -> Self::T {
878 Path {
879 a: MInt::one(),
880 b: MInt::zero(),
881 sum: MInt::zero(),
882 cnt: MInt::zero(),
883 }
884 }
885 }
886 impl Associative for PathMonoid {}
887
888 struct PathPairMonoid;
889 impl Magma for PathPairMonoid {
890 type T = PathPair;
891 fn operate(x: &Self::T, y: &Self::T) -> Self::T {
892 PathPair {
893 forward: PathMonoid::operate(&x.forward, &y.forward),
894 reverse: PathMonoid::operate(&y.reverse, &x.reverse),
895 }
896 }
897 }
898 impl Unital for PathPairMonoid {
899 fn unit() -> Self::T {
900 PathPair {
901 forward: PathMonoid::unit(),
902 reverse: PathMonoid::unit(),
903 }
904 }
905 }
906 impl Associative for PathPairMonoid {}
907
908 struct FixedCluster;
909 impl MonoidCluster for FixedCluster {
910 type Vertex = MInt;
911 type Edge = (MInt, MInt);
912 type PointMonoid = PointMonoid;
913 type PathMonoid = PathMonoid;
914
915 fn add_vertex(point: &Point, vertex: &MInt, parent_edge: Option<&(MInt, MInt)>) -> Path {
916 let cnt = point.cnt + MInt::one();
917 let subtotal = point.sum + *vertex;
918 let (a, b) = parent_edge.copied().unwrap_or((MInt::one(), MInt::zero()));
919 Path {
920 a,
921 b,
922 sum: a * subtotal + b * cnt,
923 cnt,
924 }
925 }
926
927 fn add_edge(path: &Path) -> Point {
928 Point {
929 sum: path.sum,
930 cnt: path.cnt,
931 }
932 }
933 }
934
935 struct RerootCluster;
936 impl MonoidCluster for RerootCluster {
937 type Vertex = MInt;
938 type Edge = (MInt, MInt);
939 type PointMonoid = PointMonoid;
940 type PathMonoid = PathPairMonoid;
941
942 fn add_vertex(
943 point: &Point,
944 vertex: &MInt,
945 parent_edge: Option<&(MInt, MInt)>,
946 ) -> PathPair {
947 let cnt = point.cnt + MInt::one();
948 let subtotal = point.sum + *vertex;
949 let (a, b) = parent_edge.copied().unwrap_or((MInt::one(), MInt::zero()));
950 PathPair {
951 forward: Path {
952 a,
953 b,
954 sum: a * subtotal + b * cnt,
955 cnt,
956 },
957 reverse: Path {
958 a,
959 b,
960 sum: subtotal,
961 cnt,
962 },
963 }
964 }
965
966 fn add_edge(path: &PathPair) -> Point {
967 Point {
968 sum: path.forward.sum,
969 cnt: path.forward.cnt,
970 }
971 }
972 }
973
974 const MATCHING_NEG_INF: i64 = i64::MIN / 4;
975
976 struct MatchingCluster;
977
978 impl MatchingCluster {
979 fn point() -> [Vec<i64>; 2] {
980 array::from_fn(|_| Vec::new())
981 }
982
983 fn path() -> [[Vec<i64>; 2]; 2] {
984 array::from_fn(|_| Self::point())
985 }
986
987 fn relax(dst: &mut Vec<i64>, src: &[i64], shift: usize, add: i64) {
988 for (i, &x) in src.iter().enumerate() {
989 if x == MATCHING_NEG_INF {
990 continue;
991 }
992 let j = i + shift;
993 if dst.len() <= j {
994 dst.resize(j + 1, MATCHING_NEG_INF);
995 }
996 dst[j] = dst[j].max(x + add);
997 }
998 }
999
1000 fn convolve(dst: &mut Vec<i64>, a: &[i64], b: &[i64]) {
1001 if a.is_empty() || b.is_empty() {
1002 return;
1003 }
1004 let len = a.len() + b.len() - 1;
1005 if dst.len() < len {
1006 dst.resize(len, MATCHING_NEG_INF);
1007 }
1008 for (i, &x) in a.iter().enumerate() {
1009 if x == MATCHING_NEG_INF {
1010 continue;
1011 }
1012 for (j, &y) in b.iter().enumerate() {
1013 if y != MATCHING_NEG_INF {
1014 dst[i + j] = dst[i + j].max(x + y);
1015 }
1016 }
1017 }
1018 }
1019 }
1020
1021 impl Cluster for MatchingCluster {
1022 type Vertex = ();
1023 type Edge = i64;
1024 type Path = [[Vec<i64>; 2]; 2];
1025 type Point = [Vec<i64>; 2];
1026
1027 fn unit_path() -> Self::Path {
1028 let mut path = Self::path();
1029 path[0][0] = vec![0];
1030 path
1031 }
1032
1033 fn unit_point() -> Self::Point {
1034 let mut point = Self::point();
1035 point[0] = vec![0];
1036 point
1037 }
1038
1039 fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path {
1040 let mut path = Self::path();
1041 for (lt, left_row) in left.iter().enumerate() {
1042 for (lb, left_seq) in left_row.iter().enumerate() {
1043 for (rt, right_row) in right.iter().enumerate() {
1044 if lb == 1 && rt == 1 {
1045 continue;
1046 }
1047 for (rb, right_seq) in right_row.iter().enumerate() {
1048 Self::convolve(&mut path[lt][rb], left_seq, right_seq);
1049 }
1050 }
1051 }
1052 }
1053 path
1054 }
1055
1056 fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point {
1057 let mut point = Self::point();
1058 for (lu, left_seq) in left.iter().enumerate() {
1059 for (ru, right_seq) in right.iter().enumerate() {
1060 if lu == 1 && ru == 1 {
1061 continue;
1062 }
1063 Self::convolve(&mut point[lu | ru], left_seq, right_seq);
1064 }
1065 }
1066 point
1067 }
1068
1069 fn add_edge(path: &Self::Path) -> Self::Point {
1070 let mut point = Self::point();
1071 for (top, row) in path.iter().enumerate() {
1072 for seq in row {
1073 Self::relax(&mut point[top], seq, 0, 0);
1074 }
1075 }
1076 point
1077 }
1078
1079 fn add_vertex(
1080 point: &Self::Point,
1081 _vertex: &Self::Vertex,
1082 parent_edge: Option<&Self::Edge>,
1083 ) -> Self::Path {
1084 let mut path = Self::path();
1085 for (used, seq) in point.iter().enumerate() {
1086 Self::relax(&mut path[0][used], seq, 0, 0);
1087 }
1088 if let Some(&w) = parent_edge {
1089 Self::relax(&mut path[1][1], &point[0], 1, w);
1090 }
1091 path
1092 }
1093 }
1094
1095 fn naive_rooted(
1096 graph: &UndirectedSparseGraph,
1097 vertices: &[MInt],
1098 edges: &[(MInt, MInt)],
1099 root: usize,
1100 ) -> Point {
1101 fn dfs(
1102 graph: &UndirectedSparseGraph,
1103 vertices: &[MInt],
1104 edges: &[(MInt, MInt)],
1105 u: usize,
1106 p: usize,
1107 in_edge: Option<usize>,
1108 ) -> Point {
1109 let mut point = PointMonoid::unit();
1110 for a in graph.neighbors(u) {
1111 if a.to != p {
1112 point = PointMonoid::operate(
1113 &point,
1114 &dfs(graph, vertices, edges, a.to, u, Some(a.label)),
1115 );
1116 }
1117 }
1118 let cnt = point.cnt + MInt::one();
1119 let subtotal = point.sum + vertices[u];
1120 let (a, b) = in_edge
1121 .map(|eid| edges[eid])
1122 .unwrap_or((MInt::one(), MInt::zero()));
1123 Point {
1124 sum: a * subtotal + b * cnt,
1125 cnt,
1126 }
1127 }
1128 dfs(graph, vertices, edges, root, usize::MAX, None)
1129 }
1130
1131 fn gen_weights(rng: &mut Xorshift, n: usize, m: usize) -> (Vec<MInt>, Vec<(MInt, MInt)>) {
1132 let vertices = (0..n)
1133 .map(|_| MInt::from(rng.random(0u32..10)))
1134 .collect::<Vec<_>>();
1135 let edges = (0..m)
1136 .map(|_| {
1137 (
1138 MInt::from(rng.random(0u32..10)),
1139 MInt::from(rng.random(0u32..10)),
1140 )
1141 })
1142 .collect::<Vec<_>>();
1143 (vertices, edges)
1144 }
1145
1146 fn run_fixed_case(graph: &UndirectedSparseGraph, rounds: usize, rng: &mut Xorshift) {
1147 let n = graph.vertices_size();
1148 let m = graph.edges_size();
1149 let (mut vertices, mut edges) = gen_weights(rng, n, m);
1150 let tree = graph.static_top_tree(0);
1151 assert_eq!(
1152 tree.fold_all::<FixedCluster>(&vertices, &edges),
1153 naive_rooted(graph, &vertices, &edges, 0)
1154 );
1155 let mut dp = tree.dp::<FixedCluster>(vertices.clone(), edges.clone());
1156 assert_eq!(*dp.fold_all(), naive_rooted(graph, &vertices, &edges, 0));
1157 for (v, &vertex) in vertices.iter().enumerate() {
1158 assert_eq!(*dp.get_vertex(v), vertex);
1159 }
1160 for (eid, &edge) in edges.iter().enumerate() {
1161 assert_eq!(*dp.get_edge(eid), edge);
1162 }
1163
1164 for _ in 0..rounds {
1165 if rng.random(0u32..2) == 0 {
1166 let v = rng.random(0..n);
1167 let x = MInt::from(rng.random(0u32..20));
1168 vertices[v] += x;
1169 dp.apply_vertex(v, |value| *value += x);
1170 assert_eq!(*dp.get_vertex(v), vertices[v]);
1171 } else if m > 0 {
1172 let eid = rng.random(0..m);
1173 let edge = (
1174 MInt::from(rng.random(0u32..20)),
1175 MInt::from(rng.random(0u32..20)),
1176 );
1177 edges[eid].0 += edge.0;
1178 edges[eid].1 += edge.1;
1179 dp.apply_edge(eid, |value| {
1180 value.0 += edge.0;
1181 value.1 += edge.1;
1182 });
1183 assert_eq!(*dp.get_edge(eid), edges[eid]);
1184 }
1185 assert_eq!(*dp.fold_all(), naive_rooted(graph, &vertices, &edges, 0));
1186 }
1187 }
1188
1189 fn run_reroot_case(graph: &UndirectedSparseGraph, rounds: usize, rng: &mut Xorshift) {
1190 let n = graph.vertices_size();
1191 let m = graph.edges_size();
1192 let (mut vertices, mut edges) = gen_weights(rng, n, m);
1193 let tree = graph.static_top_tree(0);
1194 let mut dp = tree.dp::<RerootCluster>(vertices.clone(), edges.clone());
1195 assert_eq!(
1196 dp.fold_all().sum,
1197 naive_rooted(graph, &vertices, &edges, 0).sum
1198 );
1199
1200 for _ in 0..rounds {
1201 if rng.random(0u32..2) == 0 {
1202 let v = rng.random(0..n);
1203 let x = MInt::from(rng.random(0u32..20));
1204 vertices[v] = x;
1205 dp.set_vertex(v, x);
1206 } else if m > 0 {
1207 let eid = rng.random(0..m);
1208 let edge = (
1209 MInt::from(rng.random(0u32..20)),
1210 MInt::from(rng.random(0u32..20)),
1211 );
1212 edges[eid] = edge;
1213 dp.set_edge(eid, edge);
1214 } else {
1215 let v = rng.random(0..n);
1216 let x = MInt::from(rng.random(0u32..20));
1217 vertices[v] = x;
1218 dp.set_vertex(v, x);
1219 }
1220 for root in 0..n {
1221 let got = dp.fold_path(root).reverse.sum;
1222 let want = naive_rooted(graph, &vertices, &edges, root).sum;
1223 assert_eq!(got, want, "root={root}");
1224 }
1225 }
1226 }
1227
1228 #[test]
1229 fn static_top_tree_fixed_random() {
1230 let mut rng = Xorshift::default();
1231 for n in 1..=16 {
1232 for graph in [
1233 rng.random(PathTree(n)),
1234 rng.random(StarTree(n)),
1235 UndirectedSparseGraph::from_edges(n, (1..n).map(|v| ((v - 1) / 2, v)).collect()),
1236 ] {
1237 run_fixed_case(&graph, 30, &mut rng);
1238 }
1239 }
1240 for _ in 0..30 {
1241 let graph = rng.random(MixedTree(1..=14usize));
1242 run_fixed_case(&graph, 40, &mut rng);
1243 }
1244 }
1245
1246 #[test]
1247 fn static_top_tree_reroot_random() {
1248 let mut rng = Xorshift::default();
1249 for n in 1..=16 {
1250 for graph in [
1251 rng.random(PathTree(n)),
1252 rng.random(StarTree(n)),
1253 UndirectedSparseGraph::from_edges(n, (1..n).map(|v| ((v - 1) / 2, v)).collect()),
1254 ] {
1255 run_reroot_case(&graph, 30, &mut rng);
1256 }
1257 }
1258 for _ in 0..20 {
1259 let graph = rng.random(MixedTree(1..=12usize));
1260 run_reroot_case(&graph, 30, &mut rng);
1261 }
1262 }
1263
1264 #[test]
1265 fn static_top_tree_matching_matches_bruteforce() {
1266 fn brute_matching(graph: &UndirectedSparseGraph, weights: &[i64]) -> Vec<i64> {
1267 let mut best = vec![MATCHING_NEG_INF; graph.vertices_size()];
1268 best[0] = 0;
1269 for mask in 0usize..1usize << graph.edges_size() {
1270 let mut used = vec![false; graph.vertices_size()];
1271 let mut count = 0usize;
1272 let mut sum = 0i64;
1273 let mut ok = true;
1274 for (eid, &(u, v)) in graph.edges.iter().enumerate() {
1275 if mask >> eid & 1 == 0 {
1276 continue;
1277 }
1278 if used[u] || used[v] {
1279 ok = false;
1280 break;
1281 }
1282 used[u] = true;
1283 used[v] = true;
1284 count += 1;
1285 sum += weights[eid];
1286 }
1287 if ok {
1288 best[count] = best[count].max(sum);
1289 }
1290 }
1291 best
1292 }
1293
1294 let mut rng = Xorshift::default();
1295 let mut graphs = Vec::new();
1296 for n in 1..=11 {
1297 graphs.extend([
1298 rng.random(PathTree(n)),
1299 rng.random(StarTree(n)),
1300 UndirectedSparseGraph::from_edges(n, (1..n).map(|v| ((v - 1) / 2, v)).collect()),
1301 ]);
1302 }
1303 graphs.extend((0..100).map(|_| rng.random(MixedTree(1..=11usize))));
1304 for (case, graph) in graphs.into_iter().enumerate() {
1305 let weights = rng
1306 .random_iter(-20i64..=20)
1307 .take(graph.edges_size())
1308 .collect::<Vec<_>>();
1309 let tree = graph.static_top_tree(0);
1310 let mut got = tree
1311 .fold_all::<MatchingCluster>(&vec![(); graph.vertices_size()], &weights)[0]
1312 .clone();
1313 got.resize(graph.vertices_size(), MATCHING_NEG_INF);
1314 assert_eq!(got, brute_matching(&graph, &weights), "case={case}");
1315 }
1316 }
1317}