competitive/tree/
top_tree.rs

1use super::{
2    Allocator, Magma, MemoryPool, Monoid, Unital,
3    binary_search_tree::{
4        BstDataMutRef, BstNode, BstRoot, BstSeeker, BstSpec, EqualSide, node::WithParent,
5    },
6    splay_operations,
7};
8use std::{marker::PhantomData, mem::replace, ptr::NonNull};
9
10/// `compress` must be associative and `rake` must be associative and
11/// commutative. `reverse` must be an involution and reverse the operand order
12/// of `compress`.
13pub trait TopTreeSpec: Sized {
14    type Info;
15    type Point: Clone;
16    type Path: Clone;
17
18    fn vertex(info: &Self::Info) -> Self::Path;
19    fn add_vertex(point: &Self::Point, info: &Self::Info) -> Self::Path;
20    fn add_edge(path: &Self::Path) -> Self::Point;
21    fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point;
22    fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path;
23    fn reverse(path: &mut Self::Path);
24}
25
26/// `act_path` applies the heavy contribution and `act_path_light` applies the
27/// light contribution represented in the same path aggregate.
28pub trait TopTreeAction<S>
29where
30    S: TopTreeSpec,
31{
32    type Action: PartialEq;
33    /// Composes actions in application order: `operate(a, b)` applies `a`
34    /// before `b`.
35    type ActionMonoid: Monoid<T = Self::Action>;
36
37    /// Whether pending actions require propagation from the auxiliary root before rotations.
38    const ROOT_TO_NODE_TOP_DOWN: bool = true;
39
40    fn act_info(info: &mut S::Info, action: &Self::Action);
41    fn act_point(point: &mut S::Point, action: &Self::Action);
42    fn act_path(path: &mut S::Path, action: &Self::Action);
43    fn act_path_light(path: &mut S::Path, action: &Self::Action);
44}
45
46pub struct NoTopTreeAction;
47
48impl<S> TopTreeAction<S> for NoTopTreeAction
49where
50    S: TopTreeSpec,
51{
52    type Action = ();
53    type ActionMonoid = ();
54
55    const ROOT_TO_NODE_TOP_DOWN: bool = false;
56
57    fn act_info(_info: &mut S::Info, _action: &Self::Action) {}
58    fn act_point(_point: &mut S::Point, _action: &Self::Action) {}
59    fn act_path(_path: &mut S::Path, _action: &Self::Action) {}
60    fn act_path_light(_path: &mut S::Path, _action: &Self::Action) {}
61}
62
63struct TopTreeData<S, A>
64where
65    S: TopTreeSpec,
66    A: TopTreeAction<S>,
67{
68    info: S::Info,
69    sum: S::Path,
70    light: Option<RakePtr<S, A>>,
71    belong: Option<RakePtr<S, A>>,
72    heavy_action: A::Action,
73    light_action: A::Action,
74    index_and_reverse: usize,
75}
76
77struct RakeData<S, A>
78where
79    S: TopTreeSpec,
80    A: TopTreeAction<S>,
81{
82    key: S::Point,
83    sum: S::Point,
84    action: A::Action,
85    buffer: A::Action,
86}
87
88struct TopBstSpec<S, A>(PhantomData<fn() -> (S, A)>);
89struct RakeBstSpec<S, A>(PhantomData<fn() -> (S, A)>);
90
91type TopNode<S, A> = BstNode<TopTreeData<S, A>, WithParent<TopTreeData<S, A>>>;
92type RakeNode<S, A> = BstNode<RakeData<S, A>, WithParent<RakeData<S, A>>>;
93type TopPtr<S, A> = NonNull<TopNode<S, A>>;
94type RakePtr<S, A> = NonNull<RakeNode<S, A>>;
95
96impl<S, A> TopBstSpec<S, A>
97where
98    S: TopTreeSpec,
99    A: TopTreeAction<S>,
100{
101    #[inline]
102    fn is_unit(action: &A::Action) -> bool {
103        <A::ActionMonoid as Unital>::is_unit(action)
104    }
105
106    #[inline]
107    fn compose(target: &mut A::Action, action: &A::Action) {
108        <A::ActionMonoid as Magma>::operate_assign(target, action);
109    }
110
111    #[inline]
112    unsafe fn toggle(mut node: TopPtr<S, A>) {
113        unsafe { node.as_mut().child.swap(0, 1) };
114        let data = unsafe { &mut node.as_mut().data };
115        data.index_and_reverse ^= 1;
116        S::reverse(&mut data.sum);
117    }
118
119    #[inline]
120    unsafe fn apply_heavy(mut node: TopPtr<S, A>, action: &A::Action) {
121        let data = unsafe { &mut node.as_mut().data };
122        Self::compose(&mut data.heavy_action, action);
123        A::act_info(&mut data.info, action);
124        A::act_path(&mut data.sum, action);
125    }
126
127    #[inline]
128    unsafe fn apply_light(mut node: TopPtr<S, A>, action: &A::Action) {
129        let data = unsafe { &mut node.as_mut().data };
130        Self::compose(&mut data.light_action, action);
131        A::act_path_light(&mut data.sum, action);
132    }
133
134    #[inline]
135    unsafe fn apply_all(mut node: TopPtr<S, A>, action: &A::Action) {
136        if Self::is_unit(action) {
137            return;
138        }
139        let data = unsafe { &mut node.as_mut().data };
140        Self::compose(&mut data.heavy_action, action);
141        Self::compose(&mut data.light_action, action);
142        A::act_info(&mut data.info, action);
143        A::act_path(&mut data.sum, action);
144        A::act_path_light(&mut data.sum, action);
145    }
146}
147
148impl<S, A> BstSpec for TopBstSpec<S, A>
149where
150    S: TopTreeSpec,
151    A: TopTreeAction<S>,
152{
153    type Parent = WithParent<Self::Data>;
154    type Data = TopTreeData<S, A>;
155
156    #[inline]
157    fn top_down(mut node: BstDataMutRef<'_, Self>) {
158        let pointer = node.node;
159        if node.reborrow().into_data().index_and_reverse & 1 != 0 {
160            node.data_mut().index_and_reverse &= !1;
161            for child in unsafe { pointer.as_ref().child }.into_iter().flatten() {
162                unsafe { Self::toggle(child) };
163            }
164        }
165
166        if !Self::is_unit(&node.reborrow().into_data().heavy_action) {
167            let action = replace(
168                &mut node.data_mut().heavy_action,
169                <A::ActionMonoid as Unital>::unit(),
170            );
171            for child in unsafe { pointer.as_ref().child }.into_iter().flatten() {
172                unsafe { Self::apply_heavy(child, &action) };
173            }
174        }
175        if !Self::is_unit(&node.reborrow().into_data().light_action) {
176            let action = replace(
177                &mut node.data_mut().light_action,
178                <A::ActionMonoid as Unital>::unit(),
179            );
180            for child in unsafe { pointer.as_ref().child }.into_iter().flatten() {
181                unsafe { Self::apply_light(child, &action) };
182            }
183            if let Some(light) = node.reborrow().into_data().light {
184                unsafe { RakeBstSpec::<S, A>::apply(light, &action) };
185            }
186        }
187    }
188
189    #[inline]
190    fn bottom_up(node: BstDataMutRef<'_, Self>) {
191        let pointer = node.node;
192        let data = unsafe { &mut (*pointer.as_ptr()).data };
193        let mut sum = if let Some(light) = data.light {
194            S::add_vertex(unsafe { &light.as_ref().data.sum }, &data.info)
195        } else {
196            S::vertex(&data.info)
197        };
198        if let Some(left) = unsafe { pointer.as_ref().child[0] } {
199            sum = S::compress(unsafe { &left.as_ref().data.sum }, &sum);
200        }
201        if let Some(right) = unsafe { pointer.as_ref().child[1] } {
202            sum = S::compress(&sum, unsafe { &right.as_ref().data.sum });
203        }
204        data.sum = sum;
205    }
206
207    fn merge(_left: Option<BstRoot<Self>>, _right: Option<BstRoot<Self>>) -> Option<BstRoot<Self>> {
208        unreachable!("top trees do not merge auxiliary trees through BstSpec")
209    }
210
211    fn split<Seeker>(
212        _node: Option<BstRoot<Self>>,
213        _seeker: Seeker,
214        _equal_side: EqualSide,
215    ) -> (Option<BstRoot<Self>>, Option<BstRoot<Self>>)
216    where
217        Seeker: BstSeeker<Spec = Self>,
218    {
219        unreachable!("top trees do not split auxiliary trees through BstSpec")
220    }
221}
222
223impl<S, A> RakeBstSpec<S, A>
224where
225    S: TopTreeSpec,
226    A: TopTreeAction<S>,
227{
228    #[inline]
229    unsafe fn apply(mut node: RakePtr<S, A>, action: &A::Action) {
230        let data = unsafe { &mut node.as_mut().data };
231        A::act_point(&mut data.key, action);
232        A::act_point(&mut data.sum, action);
233        TopBstSpec::<S, A>::compose(&mut data.action, action);
234        TopBstSpec::<S, A>::compose(&mut data.buffer, action);
235    }
236}
237
238impl<S, A> BstSpec for RakeBstSpec<S, A>
239where
240    S: TopTreeSpec,
241    A: TopTreeAction<S>,
242{
243    type Parent = WithParent<Self::Data>;
244    type Data = RakeData<S, A>;
245
246    #[inline]
247    fn top_down(mut node: BstDataMutRef<'_, Self>) {
248        if TopBstSpec::<S, A>::is_unit(&node.reborrow().into_data().action) {
249            return;
250        }
251        let pointer = node.node;
252        let action = replace(
253            &mut node.data_mut().action,
254            <A::ActionMonoid as Unital>::unit(),
255        );
256        for child in unsafe { pointer.as_ref().child }.into_iter().flatten() {
257            unsafe { Self::apply(child, &action) };
258        }
259    }
260
261    #[inline]
262    fn bottom_up(node: BstDataMutRef<'_, Self>) {
263        let pointer = node.node;
264        let data = unsafe { &mut (*pointer.as_ptr()).data };
265        let mut sum = data.key.clone();
266        if let Some(left) = unsafe { pointer.as_ref().child[0] } {
267            sum = S::rake(&sum, unsafe { &left.as_ref().data.sum });
268        }
269        if let Some(right) = unsafe { pointer.as_ref().child[1] } {
270            sum = S::rake(&sum, unsafe { &right.as_ref().data.sum });
271        }
272        data.sum = sum;
273    }
274
275    fn merge(_left: Option<BstRoot<Self>>, _right: Option<BstRoot<Self>>) -> Option<BstRoot<Self>> {
276        unreachable!("rake trees use their dedicated merge operation")
277    }
278
279    fn split<Seeker>(
280        _node: Option<BstRoot<Self>>,
281        _seeker: Seeker,
282        _equal_side: EqualSide,
283    ) -> (Option<BstRoot<Self>>, Option<BstRoot<Self>>)
284    where
285        Seeker: BstSeeker<Spec = Self>,
286    {
287        unreachable!("rake trees do not use BstSpec::split")
288    }
289}
290
291/// A self-adjusting top tree, also called a strong link-cut tree.
292///
293/// This is not a classical worst-case-balanced top tree. Circular order and
294/// `select` are not supported.
295pub struct TopTree<S, A = NoTopTreeAction>
296where
297    S: TopTreeSpec,
298    A: TopTreeAction<S>,
299{
300    nodes: Vec<TopPtr<S, A>>,
301    node_allocator: MemoryPool<TopNode<S, A>>,
302    rake_allocator: MemoryPool<RakeNode<S, A>>,
303}
304
305impl<S, A> TopTree<S, A>
306where
307    S: TopTreeSpec,
308    A: TopTreeAction<S>,
309{
310    pub fn with_capacity(capacity: usize) -> Self {
311        Self {
312            nodes: Vec::with_capacity(capacity),
313            node_allocator: MemoryPool::with_capacity(capacity),
314            rake_allocator: MemoryPool::with_capacity(capacity),
315        }
316    }
317
318    /// `edges` must form a tree over the values in iteration order.
319    pub fn from_edges<T>(values: T, edges: &[(usize, usize)]) -> Self
320    where
321        T: IntoIterator<Item = S::Info>,
322    {
323        let mut tree: Self = values.into_iter().collect();
324        for (child, parent, preferred) in
325            splay_operations::rooted_heavy_order(tree.nodes.len(), edges)
326                .into_iter()
327                .rev()
328        {
329            let child = tree.node(child);
330            let mut parent = tree.node(parent);
331            unsafe {
332                (*child.as_ptr()).parent.parent = Some(parent);
333                if preferred {
334                    parent.as_mut().child[1] = Some(child);
335                } else {
336                    let point = S::add_edge(&child.as_ref().data.sum);
337                    let (light, entry) = tree.rake_insert(parent.as_ref().data.light, point);
338                    parent.as_mut().data.light = Some(light);
339                    (*child.as_ptr()).data.belong = Some(entry);
340                }
341                Self::pull_top(parent);
342            }
343        }
344        tree
345    }
346
347    pub fn add_node(&mut self, info: S::Info) -> usize {
348        let index = self.nodes.len();
349        let sum = S::vertex(&info);
350        let node = self.node_allocator.allocate(BstNode::new(TopTreeData {
351            info,
352            sum,
353            light: None,
354            belong: None,
355            heavy_action: <A::ActionMonoid as Unital>::unit(),
356            light_action: <A::ActionMonoid as Unital>::unit(),
357            index_and_reverse: index << 1,
358        }));
359        self.nodes.push(node);
360        index
361    }
362
363    fn node(&self, index: usize) -> TopPtr<S, A> {
364        self.nodes[index]
365    }
366
367    #[inline]
368    unsafe fn pull_top(node: TopPtr<S, A>) {
369        unsafe { TopBstSpec::<S, A>::bottom_up(BstDataMutRef::new_unchecked(node)) };
370    }
371
372    #[inline]
373    unsafe fn pull_rake(node: RakePtr<S, A>) {
374        unsafe { RakeBstSpec::<S, A>::bottom_up(BstDataMutRef::new_unchecked(node)) };
375    }
376
377    #[inline]
378    unsafe fn splay_top(node: TopPtr<S, A>) {
379        let root = if A::ROOT_TO_NODE_TOP_DOWN {
380            unsafe {
381                splay_operations::with_parent::splay::<TopBstSpec<S, A>, TopTreeData<S, A>>(node)
382            }
383        } else {
384            unsafe {
385                splay_operations::with_parent::splay_with_local_top_down::<
386                    TopBstSpec<S, A>,
387                    TopTreeData<S, A>,
388                >(node)
389            }
390        };
391        if root != node {
392            unsafe {
393                (*node.as_ptr()).data.belong = (*root.as_ptr()).data.belong.take();
394            }
395        }
396    }
397
398    #[inline]
399    unsafe fn splay_rake(node: RakePtr<S, A>) {
400        unsafe {
401            splay_operations::with_parent::splay_with_local_top_down::<
402                RakeBstSpec<S, A>,
403                RakeData<S, A>,
404            >(node)
405        };
406    }
407
408    unsafe fn rake_rightmost(mut node: RakePtr<S, A>) -> RakePtr<S, A> {
409        loop {
410            unsafe { RakeBstSpec::<S, A>::top_down(BstDataMutRef::new_unchecked(node)) };
411            match unsafe { node.as_ref().child[1] } {
412                Some(right) => node = right,
413                None => return node,
414            }
415        }
416    }
417
418    unsafe fn rake_insert(
419        &mut self,
420        root: Option<RakePtr<S, A>>,
421        key: S::Point,
422    ) -> (RakePtr<S, A>, RakePtr<S, A>) {
423        let mut node = self.rake_allocator.allocate(BstNode::new(RakeData {
424            sum: key.clone(),
425            key,
426            action: <A::ActionMonoid as Unital>::unit(),
427            buffer: <A::ActionMonoid as Unital>::unit(),
428        }));
429        if let Some(mut root) = root {
430            unsafe {
431                node.as_mut().child[0] = Some(root);
432                root.as_mut().parent.parent = Some(node);
433                Self::pull_rake(node);
434            }
435        }
436        (node, node)
437    }
438
439    unsafe fn rake_remove(
440        &mut self,
441        mut node: RakePtr<S, A>,
442    ) -> (Option<RakePtr<S, A>>, A::Action) {
443        unsafe {
444            Self::splay_rake(node);
445            RakeBstSpec::<S, A>::top_down(BstDataMutRef::new_unchecked(node));
446        }
447        let left = unsafe { node.as_mut().child[0].take() };
448        let right = unsafe { node.as_mut().child[1].take() };
449        for mut child in [left, right].into_iter().flatten() {
450            unsafe { child.as_mut().parent.parent = None };
451        }
452        let root = match (left, right) {
453            (None, right) => right,
454            (left, None) => left,
455            (Some(left), Some(right)) => {
456                let mut root = unsafe { Self::rake_rightmost(left) };
457                unsafe {
458                    Self::splay_rake(root);
459                    root.as_mut().child[1] = Some(right);
460                    (*right.as_ptr()).parent.parent = Some(root);
461                    Self::pull_rake(root);
462                }
463                Some(root)
464            }
465        };
466        let node = self.rake_allocator.deallocate(node);
467        (root, node.data.buffer)
468    }
469
470    fn access_node(&mut self, node: TopPtr<S, A>) {
471        unsafe {
472            let mut previous: Option<TopPtr<S, A>> = None;
473            let mut current = Some(node);
474            while let Some(mut cursor) = current {
475                Self::splay_top(cursor);
476                let next = cursor.as_ref().parent.parent;
477                if let Some(right) = cursor.as_mut().child[1].take() {
478                    let point = S::add_edge(&right.as_ref().data.sum);
479                    let (light, entry) = self.rake_insert(cursor.as_ref().data.light, point);
480                    cursor.as_mut().data.light = Some(light);
481                    (*right.as_ptr()).data.belong = Some(entry);
482                }
483                if let Some(previous) = previous {
484                    let entry = (*previous.as_ptr())
485                        .data
486                        .belong
487                        .take()
488                        .expect("a virtual path must have a rake-tree entry");
489                    let (light, action) = self.rake_remove(entry);
490                    cursor.as_mut().data.light = light;
491                    TopBstSpec::<S, A>::apply_all(previous, &action);
492                    cursor.as_mut().child[1] = Some(previous);
493                    (*previous.as_ptr()).parent.parent = Some(cursor);
494                }
495                Self::pull_top(cursor);
496                previous = Some(cursor);
497                current = next;
498            }
499            Self::splay_top(node);
500        }
501    }
502
503    pub fn get(&mut self, node: usize) -> &S::Info {
504        let node = self.node(node);
505        self.access_node(node);
506        unsafe { &node.as_ref().data.info }
507    }
508
509    pub fn set(&mut self, node: usize, info: S::Info) {
510        self.modify(node, |_| info);
511    }
512
513    pub fn modify<F>(&mut self, node: usize, f: F)
514    where
515        F: FnOnce(&S::Info) -> S::Info,
516    {
517        let mut node = self.node(node);
518        self.access_node(node);
519        unsafe {
520            node.as_mut().data.info = f(&node.as_ref().data.info);
521            Self::pull_top(node);
522        }
523    }
524
525    pub fn reroot(&mut self, node: usize) {
526        let node = self.node(node);
527        self.access_node(node);
528        unsafe { TopBstSpec::<S, A>::toggle(node) };
529    }
530
531    /// `child` and `parent` must belong to different trees.
532    pub fn link(&mut self, child: usize, parent: usize) {
533        assert_ne!(child, parent);
534        self.reroot(child);
535        let child = self.node(child);
536        let mut parent = self.node(parent);
537        self.access_node(parent);
538        unsafe {
539            (*child.as_ptr()).parent.parent = Some(parent);
540            let point = S::add_edge(&child.as_ref().data.sum);
541            let (light, entry) = self.rake_insert(parent.as_ref().data.light, point);
542            parent.as_mut().data.light = Some(light);
543            (*child.as_ptr()).data.belong = Some(entry);
544            Self::pull_top(parent);
545        }
546    }
547
548    /// `(u, v)` must be an edge.
549    pub fn cut(&mut self, u: usize, v: usize) {
550        assert_ne!(u, v);
551        self.reroot(u);
552        let mut v = self.node(v);
553        self.access_node(v);
554        unsafe {
555            let mut left = v.as_mut().child[0]
556                .take()
557                .expect("the specified edge must exist");
558            left.as_mut().parent.parent = None;
559            Self::pull_top(v);
560        }
561    }
562
563    pub fn root(&mut self, node: usize) -> usize {
564        let mut root = self.node(node);
565        self.access_node(root);
566        unsafe {
567            loop {
568                TopBstSpec::<S, A>::top_down(BstDataMutRef::new_unchecked(root));
569                match root.as_ref().child[0] {
570                    Some(left) => root = left,
571                    None => break,
572                }
573            }
574            Self::splay_top(root);
575            root.as_ref().data.index_and_reverse >> 1
576        }
577    }
578
579    pub fn is_connected(&mut self, u: usize, v: usize) -> bool {
580        self.root(u) == self.root(v)
581    }
582
583    /// `u` and `v` must be connected.
584    pub fn fold_path(&mut self, u: usize, v: usize) -> S::Path {
585        self.reroot(u);
586        let v = self.node(v);
587        self.access_node(v);
588        unsafe { v.as_ref().data.sum.clone() }
589    }
590
591    /// `u` and `v` must be connected.
592    pub fn update_path(&mut self, u: usize, v: usize, action: &A::Action) {
593        self.reroot(u);
594        let v = self.node(v);
595        self.access_node(v);
596        if !TopBstSpec::<S, A>::is_unit(action) {
597            unsafe { TopBstSpec::<S, A>::apply_heavy(v, action) };
598        }
599    }
600
601    fn detach_left<R>(mut node: TopPtr<S, A>, f: impl FnOnce(TopPtr<S, A>) -> R) -> R {
602        unsafe {
603            let left = node.as_mut().child[0].take();
604            if let Some(mut left) = left {
605                left.as_mut().parent.parent = None;
606            }
607            Self::pull_top(node);
608            let result = f(node);
609            node.as_mut().child[0] = left;
610            if let Some(mut left) = left {
611                left.as_mut().parent.parent = Some(node);
612            }
613            Self::pull_top(node);
614            result
615        }
616    }
617
618    /// `(node, parent)` must be an edge.
619    pub fn fold_subtree(&mut self, node: usize, parent: usize) -> S::Path {
620        self.reroot(parent);
621        let node = self.node(node);
622        self.access_node(node);
623        Self::detach_left(node, |node| unsafe { node.as_ref().data.sum.clone() })
624    }
625
626    /// `(node, parent)` must be an edge.
627    pub fn update_subtree(&mut self, node: usize, parent: usize, action: &A::Action) {
628        self.reroot(parent);
629        let node = self.node(node);
630        self.access_node(node);
631        Self::detach_left(node, |node| unsafe {
632            TopBstSpec::<S, A>::apply_all(node, action);
633            TopBstSpec::<S, A>::top_down(BstDataMutRef::new_unchecked(node));
634        });
635    }
636}
637
638impl<S, A> FromIterator<S::Info> for TopTree<S, A>
639where
640    S: TopTreeSpec,
641    A: TopTreeAction<S>,
642{
643    fn from_iter<T: IntoIterator<Item = S::Info>>(iter: T) -> Self {
644        let iter = iter.into_iter();
645        let (lower, _) = iter.size_hint();
646        let mut tree = Self::with_capacity(lower);
647        for info in iter {
648            tree.add_node(info);
649        }
650        tree
651    }
652}
653
654#[cfg(test)]
655mod tests {
656    use super::*;
657    use crate::{
658        algebra::LastOperation,
659        graph::{Graph, UndirectedSparseGraph},
660        tools::Xorshift,
661        tree::{MixedTree, PathTree, StarTree},
662    };
663
664    fn naive_path(adjacency: &[Vec<usize>], start: usize, goal: usize) -> Vec<usize> {
665        let mut parent = vec![usize::MAX; adjacency.len()];
666        let mut stack = vec![start];
667        while let Some(u) = stack.pop() {
668            for &v in &adjacency[u] {
669                if v != parent[u] {
670                    parent[v] = u;
671                    stack.push(v);
672                }
673            }
674        }
675        let mut current = goal;
676        let mut path = Vec::new();
677        loop {
678            path.push(current);
679            if current == start {
680                break;
681            }
682            current = parent[current];
683        }
684        path.reverse();
685        path
686    }
687
688    fn naive_subtree(adjacency: &[Vec<usize>], root: usize, parent: usize) -> Vec<usize> {
689        let mut subtree = Vec::new();
690        let mut stack = vec![(root, parent)];
691        while let Some((u, parent)) = stack.pop() {
692            subtree.push(u);
693            stack.extend(
694                adjacency[u]
695                    .iter()
696                    .filter(|&&v| v != parent)
697                    .map(|&v| (v, u)),
698            );
699        }
700        subtree
701    }
702
703    fn rewire(
704        adjacency: &mut [Vec<usize>],
705        edges: &mut [(usize, usize)],
706        rng: &mut Xorshift,
707    ) -> (usize, usize, usize, usize) {
708        let edge = rng.random(0..edges.len());
709        let (u, v) = edges[edge];
710        adjacency[u].retain(|&to| to != v);
711        adjacency[v].retain(|&to| to != u);
712
713        let mut component = vec![false; adjacency.len()];
714        let mut stack = vec![u];
715        component[u] = true;
716        while let Some(x) = stack.pop() {
717            for &to in &adjacency[x] {
718                if !component[to] {
719                    component[to] = true;
720                    stack.push(to);
721                }
722            }
723        }
724        let left = (0..adjacency.len())
725            .filter(|&x| component[x])
726            .collect::<Vec<_>>();
727        let right = (0..adjacency.len())
728            .filter(|&x| !component[x])
729            .collect::<Vec<_>>();
730        let a = left[rng.random(0..left.len())];
731        let b = right[rng.random(0..right.len())];
732
733        adjacency[a].push(b);
734        adjacency[b].push(a);
735        edges[edge] = (a, b);
736        (u, v, a, b)
737    }
738
739    struct SumTopTree;
740
741    #[derive(Clone, Debug, Eq, PartialEq)]
742    struct SumPath {
743        sum: i64,
744        path_sum: i64,
745        path_size: i64,
746        size: i64,
747    }
748
749    impl TopTreeSpec for SumTopTree {
750        type Info = i64;
751        type Point = (i64, i64);
752        type Path = SumPath;
753
754        fn vertex(info: &Self::Info) -> Self::Path {
755            SumPath {
756                sum: *info,
757                path_sum: *info,
758                path_size: 1,
759                size: 1,
760            }
761        }
762
763        fn add_vertex(point: &Self::Point, info: &Self::Info) -> Self::Path {
764            SumPath {
765                sum: point.0 + *info,
766                path_sum: *info,
767                path_size: 1,
768                size: point.1 + 1,
769            }
770        }
771
772        fn add_edge(path: &Self::Path) -> Self::Point {
773            (path.sum, path.size)
774        }
775
776        fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point {
777            (left.0 + right.0, left.1 + right.1)
778        }
779
780        fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path {
781            SumPath {
782                sum: left.sum + right.sum,
783                path_sum: left.path_sum + right.path_sum,
784                path_size: left.path_size + right.path_size,
785                size: left.size + right.size,
786            }
787        }
788
789        fn reverse(_path: &mut Self::Path) {}
790    }
791
792    struct AssignAction;
793
794    impl TopTreeAction<SumTopTree> for AssignAction {
795        type Action = Option<i64>;
796        type ActionMonoid = LastOperation<i64>;
797
798        fn act_info(info: &mut i64, action: &Self::Action) {
799            if let Some(action) = action {
800                *info = *action;
801            }
802        }
803
804        fn act_point(point: &mut (i64, i64), action: &Self::Action) {
805            if let Some(action) = action {
806                point.0 = point.1 * *action;
807            }
808        }
809
810        fn act_path(path: &mut SumPath, action: &Self::Action) {
811            if let Some(action) = action {
812                let path_sum = path.path_size * *action;
813                path.sum += path_sum - path.path_sum;
814                path.path_sum = path_sum;
815            }
816        }
817
818        fn act_path_light(path: &mut SumPath, action: &Self::Action) {
819            if let Some(action) = action {
820                path.sum = path.path_sum + (path.size - path.path_size) * *action;
821            }
822        }
823    }
824
825    struct ConcatenatePath;
826
827    impl TopTreeSpec for ConcatenatePath {
828        type Info = char;
829        type Point = ();
830        type Path = String;
831
832        fn vertex(info: &Self::Info) -> Self::Path {
833            info.to_string()
834        }
835
836        fn add_vertex(_point: &Self::Point, info: &Self::Info) -> Self::Path {
837            info.to_string()
838        }
839
840        fn add_edge(_path: &Self::Path) -> Self::Point {}
841        fn rake(_left: &Self::Point, _right: &Self::Point) -> Self::Point {}
842
843        fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path {
844            let mut result = String::with_capacity(left.len() + right.len());
845            result.push_str(left);
846            result.push_str(right);
847            result
848        }
849
850        fn reverse(path: &mut Self::Path) {
851            *path = path.chars().rev().collect();
852        }
853    }
854
855    fn run_top_tree_case(graph: &UndirectedSparseGraph, rounds: usize, rng: &mut Xorshift) {
856        let n = graph.vertices_size();
857        let mut adjacency = graph
858            .vertices()
859            .map(|u| graph.neighbors(u).map(|a| a.to).collect::<Vec<_>>())
860            .collect::<Vec<_>>();
861        let mut edges = graph.edges.clone();
862        let mut values = (0..n).map(|_| rng.random(-20i64..=20)).collect::<Vec<_>>();
863        let mut tree =
864            TopTree::<SumTopTree, AssignAction>::from_edges(values.iter().copied(), &edges);
865        let root = rng.random(0..n);
866        tree.reroot(root);
867        for u in 0..n {
868            assert_eq!(tree.root(u), root);
869        }
870
871        for _ in 0..rounds {
872            match rng.random(0..if edges.is_empty() { 2 } else { 4 }) {
873                0 => {
874                    let u = rng.random(0..n);
875                    values[u] = rng.random(-20i64..=20);
876                    tree.set(u, values[u]);
877                }
878                1 => {
879                    let u = rng.random(0..n);
880                    let v = rng.random(0..n);
881                    let path = naive_path(&adjacency, u, v);
882                    let action = rng.random(-20i64..=20);
883                    for &x in &path {
884                        values[x] = action;
885                    }
886                    tree.update_path(u, v, &Some(action));
887                }
888                2 => {
889                    let (u, v, a, b) = rewire(&mut adjacency, &mut edges, rng);
890                    tree.cut(u, v);
891                    assert!(!tree.is_connected(u, v));
892                    tree.link(a, b);
893                    assert!(tree.is_connected(u, v));
894                }
895                _ => {
896                    let &(u, v) = &edges[rng.random(0..edges.len())];
897                    let (node, parent) = if rng.random(0..2) == 0 {
898                        (u, v)
899                    } else {
900                        (v, u)
901                    };
902                    let subtree = naive_subtree(&adjacency, node, parent);
903                    let action = rng.random(-20i64..=20);
904                    for &x in &subtree {
905                        values[x] = action;
906                    }
907                    tree.update_subtree(node, parent, &Some(action));
908                }
909            }
910
911            let u = rng.random(0..n);
912            let v = rng.random(0..n);
913            let path = naive_path(&adjacency, u, v);
914            assert_eq!(
915                tree.fold_path(u, v),
916                SumPath {
917                    sum: values.iter().sum(),
918                    path_sum: path.iter().map(|&x| values[x]).sum(),
919                    path_size: path.len() as i64,
920                    size: n as i64,
921                }
922            );
923
924            if !edges.is_empty() {
925                let &(u, v) = &edges[rng.random(0..edges.len())];
926                for (node, parent) in [(u, v), (v, u)] {
927                    let subtree = naive_subtree(&adjacency, node, parent);
928                    assert_eq!(
929                        tree.fold_subtree(node, parent),
930                        SumPath {
931                            sum: subtree.iter().map(|&x| values[x]).sum(),
932                            path_sum: values[node],
933                            path_size: 1,
934                            size: subtree.len() as i64,
935                        }
936                    );
937                }
938            }
939            let u = rng.random(0..n);
940            assert_eq!(*tree.get(u), values[u]);
941        }
942    }
943
944    fn run_ordered_path_case(graph: &UndirectedSparseGraph, rounds: usize, rng: &mut Xorshift) {
945        let n = graph.vertices_size();
946        let mut adjacency = graph
947            .vertices()
948            .map(|u| graph.neighbors(u).map(|a| a.to).collect::<Vec<_>>())
949            .collect::<Vec<_>>();
950        let mut edges = graph.edges.clone();
951        let mut values = (0..n)
952            .map(|_| rng.random(b'a'..=b'z') as char)
953            .collect::<Vec<_>>();
954        let mut tree = TopTree::<ConcatenatePath>::from_edges(values.iter().copied(), &edges);
955
956        for _ in 0..rounds {
957            match rng.random(0..if edges.is_empty() { 2 } else { 3 }) {
958                0 => {
959                    let u = rng.random(0..n);
960                    values[u] = rng.random(b'a'..=b'z') as char;
961                    tree.set(u, values[u]);
962                }
963                1 => {
964                    let u = rng.random(0..n);
965                    tree.reroot(u);
966                    tree.reroot(u);
967                }
968                _ => {
969                    let (u, v, a, b) = rewire(&mut adjacency, &mut edges, rng);
970                    tree.cut(u, v);
971                    tree.link(a, b);
972                }
973            }
974
975            let u = rng.random(0..n);
976            let v = rng.random(0..n);
977            assert_eq!(
978                tree.fold_path(u, v),
979                naive_path(&adjacency, u, v)
980                    .into_iter()
981                    .map(|u| values[u])
982                    .collect::<String>()
983            );
984        }
985    }
986
987    #[test]
988    fn top_tree_noncommutative_path() {
989        let mut rng = Xorshift::default();
990        for n in 1..=14 {
991            for graph in [rng.random(PathTree(n)), rng.random(StarTree(n))] {
992                run_ordered_path_case(&graph, 300, &mut rng);
993            }
994        }
995        for _ in 0..20 {
996            let graph = rng.random(MixedTree(1..=14usize));
997            run_ordered_path_case(&graph, 300, &mut rng);
998        }
999    }
1000
1001    #[test]
1002    fn top_tree_lazy() {
1003        let mut rng = Xorshift::default();
1004        for n in 1..=14 {
1005            for graph in [rng.random(PathTree(n)), rng.random(StarTree(n))] {
1006                run_top_tree_case(&graph, 300, &mut rng);
1007            }
1008        }
1009        for _ in 0..20 {
1010            let graph = rng.random(MixedTree(1..=14usize));
1011            run_top_tree_case(&graph, 300, &mut rng);
1012        }
1013    }
1014}
competitive/tree/top_tree.rs

competitive/tree/
top_tree.rs
