competitive/tree/
xor_linked_tree.rs

1use super::{MarkedScan, Scan, ScanSource};
2use std::{marker::PhantomData, ops::Range};
3
4type Marker<T> = PhantomData<fn() -> T>;
5type BuilderMarker<P, D, H, PE, EC, X, B, E> = Marker<((P, D, H), (PE, EC), (X, B, E))>;
6type ScannerMarker<U, T, P, D, H, PE, EC, X, B, E> =
7    Marker<((U, T), (P, D, H), (PE, EC), (X, B, E))>;
8
9#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
10pub enum NoParent {}
11#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
12pub enum RecordParent {}
13#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
14pub enum NoDfsPreorder {}
15#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
16pub enum RecordDfsPreorder {}
17#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
18pub enum NoDepth {}
19#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
20pub enum RecordDepth {}
21#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
22pub enum NoParentEdge {}
23#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
24pub enum RecordParentEdge {}
25#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
26pub enum NoEdgeChild {}
27#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
28pub enum RecordEdgeChild {}
29#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
30pub enum NoEIndexed {}
31#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
32pub enum EIndexed {}
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
35pub enum NoXorBottomUpOrder {}
36#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
37pub enum RecordXorBottomUpOrder {}
38
39pub trait ParentComponent {
40    type Data;
41    fn build(parent: Vec<usize>) -> Self::Data;
42}
43impl ParentComponent for NoParent {
44    type Data = ();
45    fn build(_parent: Vec<usize>) {}
46}
47impl ParentComponent for RecordParent {
48    type Data = Vec<usize>;
49    fn build(parent: Vec<usize>) -> Vec<usize> {
50        parent
51    }
52}
53
54pub trait XorBottomUpOrderBuffer {
55    type Data;
56    fn new(n: usize) -> Self::Data;
57    fn push(data: &mut Self::Data, v: usize);
58    fn as_slice(data: &Self::Data) -> &[usize];
59}
60impl XorBottomUpOrderBuffer for NoXorBottomUpOrder {
61    type Data = ();
62    fn new(_n: usize) {}
63    fn push(_data: &mut Self::Data, _v: usize) {}
64    fn as_slice(_data: &Self::Data) -> &[usize] {
65        &[]
66    }
67}
68impl XorBottomUpOrderBuffer for RecordXorBottomUpOrder {
69    type Data = Vec<usize>;
70    fn new(n: usize) -> Vec<usize> {
71        Vec::with_capacity(n.saturating_sub(1))
72    }
73    fn push(data: &mut Vec<usize>, v: usize) {
74        data.push(v);
75    }
76    fn as_slice(data: &Vec<usize>) -> &[usize] {
77        data
78    }
79}
80
81pub trait XorBottomUpOrderComponent {
82    type Data;
83}
84impl XorBottomUpOrderComponent for NoXorBottomUpOrder {
85    type Data = ();
86}
87impl XorBottomUpOrderComponent for RecordXorBottomUpOrder {
88    type Data = Vec<usize>;
89}
90
91pub trait BuildXorBottomUpOrder<B>: XorBottomUpOrderComponent
92where
93    B: XorBottomUpOrderBuffer,
94{
95    fn build(buffer: B::Data) -> Self::Data;
96}
97impl<B> BuildXorBottomUpOrder<B> for NoXorBottomUpOrder
98where
99    B: XorBottomUpOrderBuffer,
100{
101    fn build(_buffer: B::Data) {}
102}
103impl BuildXorBottomUpOrder<RecordXorBottomUpOrder> for RecordXorBottomUpOrder {
104    fn build(buffer: Vec<usize>) -> Vec<usize> {
105        buffer
106    }
107}
108
109#[derive(Debug)]
110pub struct DfsPreorder {
111    order: Vec<usize>,
112    preorder_index: Vec<usize>,
113    subtree_end: Vec<usize>,
114}
115
116pub trait DfsPreorderComponent {
117    type Data;
118    fn build(n: usize, root: usize, parent: &[usize], xor_order: &[usize]) -> Self::Data;
119}
120impl DfsPreorderComponent for NoDfsPreorder {
121    type Data = ();
122    fn build(_n: usize, _root: usize, _parent: &[usize], _xor_order: &[usize]) {}
123}
124impl DfsPreorderComponent for RecordDfsPreorder {
125    type Data = DfsPreorder;
126    fn build(n: usize, root: usize, parent: &[usize], xor_order: &[usize]) -> DfsPreorder {
127        let mut preorder_index = vec![1usize; n];
128        for &v in xor_order {
129            preorder_index[parent[v]] += preorder_index[v];
130        }
131        let mut subtree_end = vec![0usize; n];
132        if n != 0 {
133            subtree_end[root] = n;
134        }
135        for &v in xor_order.iter().rev() {
136            let p = parent[v];
137            let size = preorder_index[v];
138            let r = preorder_index[p];
139            subtree_end[v] = r;
140            preorder_index[v] = r;
141            preorder_index[p] = r - size;
142        }
143        for i in &mut preorder_index {
144            *i -= 1;
145        }
146        let mut order = vec![0usize; n];
147        for v in 0..n {
148            order[preorder_index[v]] = v;
149        }
150        DfsPreorder {
151            order,
152            preorder_index,
153            subtree_end,
154        }
155    }
156}
157
158pub trait DepthComponent {
159    type Data;
160    fn build(n: usize, root: usize, parent: &[usize], xor_order: &[usize]) -> Self::Data;
161}
162impl DepthComponent for NoDepth {
163    type Data = ();
164    fn build(_n: usize, _root: usize, _parent: &[usize], _xor_order: &[usize]) {}
165}
166impl DepthComponent for RecordDepth {
167    type Data = Vec<usize>;
168    fn build(n: usize, _root: usize, parent: &[usize], xor_order: &[usize]) -> Vec<usize> {
169        let mut depth = vec![0usize; n];
170        for &v in xor_order.iter().rev() {
171            depth[v] = depth[parent[v]] + 1;
172        }
173        depth
174    }
175}
176
177pub trait ParentEdgeComponent {
178    type Data;
179    fn build(parent_edge: Vec<usize>) -> Self::Data;
180}
181impl ParentEdgeComponent for NoParentEdge {
182    type Data = ();
183    fn build(_parent_edge: Vec<usize>) {}
184}
185impl ParentEdgeComponent for RecordParentEdge {
186    type Data = Vec<usize>;
187    fn build(parent_edge: Vec<usize>) -> Vec<usize> {
188        parent_edge
189    }
190}
191
192pub trait EdgeChildComponent {
193    type Data;
194    fn new(m: usize) -> Self::Data;
195    fn set(data: &mut Self::Data, eid: usize, child: usize);
196}
197impl EdgeChildComponent for NoEdgeChild {
198    type Data = ();
199    fn new(_m: usize) {}
200    fn set(_data: &mut Self::Data, _eid: usize, _child: usize) {}
201}
202impl EdgeChildComponent for RecordEdgeChild {
203    type Data = Vec<usize>;
204    fn new(m: usize) -> Vec<usize> {
205        vec![0usize; m]
206    }
207    fn set(data: &mut Vec<usize>, eid: usize, child: usize) {
208        data[eid] = child;
209    }
210}
211
212pub struct XorLinkedRootedTree<
213    P = NoParent,
214    D = NoDfsPreorder,
215    H = NoDepth,
216    PE = NoParentEdge,
217    EC = NoEdgeChild,
218    X = NoXorBottomUpOrder,
219> where
220    P: ParentComponent,
221    D: DfsPreorderComponent,
222    H: DepthComponent,
223    PE: ParentEdgeComponent,
224    EC: EdgeChildComponent,
225    X: XorBottomUpOrderComponent,
226{
227    n: usize,
228    root: usize,
229    parent: P::Data,
230    dfs: D::Data,
231    depth: H::Data,
232    parent_edge: PE::Data,
233    edge_child: EC::Data,
234    xor_order: X::Data,
235    _marker: Marker<(P, D, H, PE, EC, X)>,
236}
237
238pub struct XorLinkedRootedTreeBuilder<
239    P = NoParent,
240    D = NoDfsPreorder,
241    H = NoDepth,
242    PE = NoParentEdge,
243    EC = NoEdgeChild,
244    X = NoXorBottomUpOrder,
245    B = NoXorBottomUpOrder,
246    E = NoEIndexed,
247> {
248    n: usize,
249    _marker: BuilderMarker<P, D, H, PE, EC, X, B, E>,
250}
251
252impl XorLinkedRootedTree {
253    pub fn builder(n: usize) -> XorLinkedRootedTreeBuilder {
254        XorLinkedRootedTreeBuilder {
255            n,
256            _marker: PhantomData,
257        }
258    }
259}
260
261impl<P, D, H, PE, EC, X, B, E> XorLinkedRootedTreeBuilder<P, D, H, PE, EC, X, B, E> {
262    pub fn with_parent(self) -> XorLinkedRootedTreeBuilder<RecordParent, D, H, PE, EC, X, B, E> {
263        XorLinkedRootedTreeBuilder {
264            n: self.n,
265            _marker: PhantomData,
266        }
267    }
268    pub fn with_dfs_preorder(
269        self,
270    ) -> XorLinkedRootedTreeBuilder<P, RecordDfsPreorder, H, PE, EC, X, RecordXorBottomUpOrder, E>
271    {
272        XorLinkedRootedTreeBuilder {
273            n: self.n,
274            _marker: PhantomData,
275        }
276    }
277    pub fn with_depth(
278        self,
279    ) -> XorLinkedRootedTreeBuilder<P, D, RecordDepth, PE, EC, X, RecordXorBottomUpOrder, E> {
280        XorLinkedRootedTreeBuilder {
281            n: self.n,
282            _marker: PhantomData,
283        }
284    }
285    pub fn with_xor_bottom_up_order(
286        self,
287    ) -> XorLinkedRootedTreeBuilder<
288        P,
289        D,
290        H,
291        PE,
292        EC,
293        RecordXorBottomUpOrder,
294        RecordXorBottomUpOrder,
295        E,
296    > {
297        XorLinkedRootedTreeBuilder {
298            n: self.n,
299            _marker: PhantomData,
300        }
301    }
302}
303
304impl<P, D, H, X, B>
305    XorLinkedRootedTreeBuilder<P, D, H, NoParentEdge, NoEdgeChild, X, B, NoEIndexed>
306{
307    pub fn with_eindexed(
308        self,
309    ) -> XorLinkedRootedTreeBuilder<P, D, H, NoParentEdge, NoEdgeChild, X, B, EIndexed> {
310        XorLinkedRootedTreeBuilder {
311            n: self.n,
312            _marker: PhantomData,
313        }
314    }
315}
316
317impl<P, D, H, PE, EC, X, B> XorLinkedRootedTreeBuilder<P, D, H, PE, EC, X, B, EIndexed> {
318    pub fn with_parent_edge(
319        self,
320    ) -> XorLinkedRootedTreeBuilder<P, D, H, RecordParentEdge, EC, X, B, EIndexed> {
321        XorLinkedRootedTreeBuilder {
322            n: self.n,
323            _marker: PhantomData,
324        }
325    }
326    pub fn with_edge_child(
327        self,
328    ) -> XorLinkedRootedTreeBuilder<P, D, H, PE, RecordEdgeChild, X, B, EIndexed> {
329        XorLinkedRootedTreeBuilder {
330            n: self.n,
331            _marker: PhantomData,
332        }
333    }
334}
335
336impl<P, D, H, X, B> XorLinkedRootedTreeBuilder<P, D, H, NoParentEdge, NoEdgeChild, X, B, NoEIndexed>
337where
338    P: ParentComponent,
339    D: DfsPreorderComponent,
340    H: DepthComponent,
341    X: BuildXorBottomUpOrder<B>,
342    B: XorBottomUpOrderBuffer,
343{
344    /// Builds a tree rooted at 0. The `n - 1` parents are in vertex order,
345    /// and the parent of vertex `v` must be less than `v`.
346    pub fn build_from_ordered_parents(
347        self,
348        parents: impl IntoIterator<Item = usize>,
349    ) -> XorLinkedRootedTree<P, D, H, NoParentEdge, NoEdgeChild, X> {
350        let mut parent = Vec::with_capacity(self.n);
351        if self.n != 0 {
352            parent.push(usize::MAX);
353        }
354        parent.extend(parents);
355        assert_eq!(parent.len(), self.n);
356        let mut order = B::new(self.n);
357        for v in (1..self.n).rev() {
358            B::push(&mut order, v);
359        }
360        finish_rooted_tree::<P, D, H, X, B>(0, parent, order)
361    }
362
363    pub fn build<I>(
364        self,
365        root: usize,
366        edges: I,
367    ) -> XorLinkedRootedTree<P, D, H, NoParentEdge, NoEdgeChild, X>
368    where
369        I: IntoIterator<Item = (usize, usize)>,
370    {
371        let mut acc = XorAccumulator::new(self.n);
372        for (u, v) in edges {
373            acc.add_edge(u, v);
374        }
375        finish_accumulator::<P, D, H, X, B>(self.n, root, acc)
376    }
377}
378
379impl<P, D, H, PE, EC, X, B> XorLinkedRootedTreeBuilder<P, D, H, PE, EC, X, B, EIndexed>
380where
381    P: ParentComponent,
382    D: DfsPreorderComponent,
383    H: DepthComponent,
384    PE: ParentEdgeComponent,
385    EC: EdgeChildComponent,
386    X: BuildXorBottomUpOrder<B>,
387    B: XorBottomUpOrderBuffer,
388{
389    pub fn build<I>(self, root: usize, edges: I) -> XorLinkedRootedTree<P, D, H, PE, EC, X>
390    where
391        I: IntoIterator<Item = (usize, usize)>,
392    {
393        let mut acc = XorEIndexedAccumulator::new(self.n);
394        for (eid, (u, v)) in edges.into_iter().enumerate() {
395            acc.add_edge(eid, u, v);
396        }
397        finish_eindexed_accumulator::<P, D, H, PE, EC, X, B>(self.n, root, acc)
398    }
399}
400
401impl XorLinkedRootedTreeBuilder {
402    pub fn run<I, F>(self, root: usize, edges: I, f: F)
403    where
404        I: IntoIterator<Item = (usize, usize)>,
405        F: FnMut(usize, usize),
406    {
407        let mut acc = XorAccumulator::new(self.n);
408        for (u, v) in edges {
409            acc.add_edge(u, v);
410        }
411        let mut order = ();
412        acc.finish::<NoXorBottomUpOrder, _>(root, &mut order, f);
413    }
414}
415
416impl
417    XorLinkedRootedTreeBuilder<
418        NoParent,
419        NoDfsPreorder,
420        NoDepth,
421        NoParentEdge,
422        NoEdgeChild,
423        NoXorBottomUpOrder,
424        NoXorBottomUpOrder,
425        EIndexed,
426    >
427{
428    pub fn run<I, F>(self, root: usize, edges: I, f: F)
429    where
430        I: IntoIterator<Item = (usize, usize)>,
431        F: FnMut(usize, usize, usize),
432    {
433        let mut acc = XorEIndexedAccumulator::new(self.n);
434        for (eid, (u, v)) in edges.into_iter().enumerate() {
435            acc.add_edge(eid, u, v);
436        }
437        let mut order = ();
438        let mut edge_child = ();
439        acc.finish::<NoXorBottomUpOrder, NoEdgeChild, _>(root, &mut order, &mut edge_child, f);
440    }
441}
442
443impl<P, D, H, PE, EC, O> XorLinkedRootedTree<P, D, H, PE, EC, O>
444where
445    P: ParentComponent,
446    D: DfsPreorderComponent,
447    H: DepthComponent,
448    PE: ParentEdgeComponent,
449    EC: EdgeChildComponent,
450    O: XorBottomUpOrderComponent,
451{
452    pub fn vertices_size(&self) -> usize {
453        self.n
454    }
455    pub fn edges_size(&self) -> usize {
456        self.n.saturating_sub(1)
457    }
458    pub fn root(&self) -> usize {
459        self.root
460    }
461}
462
463impl<D, H, PE, EC, O> XorLinkedRootedTree<RecordParent, D, H, PE, EC, O>
464where
465    D: DfsPreorderComponent,
466    H: DepthComponent,
467    PE: ParentEdgeComponent,
468    EC: EdgeChildComponent,
469    O: XorBottomUpOrderComponent,
470{
471    pub fn parent(&self, v: usize) -> usize {
472        self.parent[v]
473    }
474    pub fn parents(&self) -> &[usize] {
475        &self.parent
476    }
477}
478
479impl<P, D, H, PE, EC> XorLinkedRootedTree<P, D, H, PE, EC, RecordXorBottomUpOrder>
480where
481    P: ParentComponent,
482    D: DfsPreorderComponent,
483    H: DepthComponent,
484    PE: ParentEdgeComponent,
485    EC: EdgeChildComponent,
486{
487    /// Returns the bottom-up XOR order, excluding the root.
488    pub fn xor_bottom_up_order(&self) -> &[usize] {
489        &self.xor_order
490    }
491    /// Returns the top-down XOR order, excluding the root.
492    pub fn xor_top_down_order(
493        &self,
494    ) -> impl DoubleEndedIterator<Item = usize> + ExactSizeIterator + '_ {
495        self.xor_order.iter().rev().copied()
496    }
497}
498
499impl<P, H, PE, EC, O> XorLinkedRootedTree<P, RecordDfsPreorder, H, PE, EC, O>
500where
501    P: ParentComponent,
502    H: DepthComponent,
503    PE: ParentEdgeComponent,
504    EC: EdgeChildComponent,
505    O: XorBottomUpOrderComponent,
506{
507    pub fn dfs_order(&self) -> &[usize] {
508        &self.dfs.order
509    }
510    pub fn dfs_index(&self, v: usize) -> usize {
511        self.dfs.preorder_index[v]
512    }
513    pub fn subtree_size(&self, v: usize) -> usize {
514        self.dfs.subtree_end[v] - self.dfs.preorder_index[v]
515    }
516    pub fn subtree_range(&self, v: usize) -> Range<usize> {
517        self.dfs.preorder_index[v]..self.dfs.subtree_end[v]
518    }
519    pub fn children(&self, v: usize) -> Children<'_> {
520        Children {
521            dfs: &self.dfs,
522            next: self.dfs.preorder_index[v] + 1,
523            end: self.dfs.subtree_end[v],
524        }
525    }
526}
527
528impl<P, D, PE, EC, O> XorLinkedRootedTree<P, D, RecordDepth, PE, EC, O>
529where
530    P: ParentComponent,
531    D: DfsPreorderComponent,
532    PE: ParentEdgeComponent,
533    EC: EdgeChildComponent,
534    O: XorBottomUpOrderComponent,
535{
536    pub fn depth(&self, v: usize) -> usize {
537        self.depth[v]
538    }
539    pub fn depths(&self) -> &[usize] {
540        &self.depth
541    }
542}
543
544impl<P, D, H, EC, O> XorLinkedRootedTree<P, D, H, RecordParentEdge, EC, O>
545where
546    P: ParentComponent,
547    D: DfsPreorderComponent,
548    H: DepthComponent,
549    EC: EdgeChildComponent,
550    O: XorBottomUpOrderComponent,
551{
552    pub fn parent_edge(&self, v: usize) -> usize {
553        self.parent_edge[v]
554    }
555    pub fn parent_edges(&self) -> &[usize] {
556        &self.parent_edge
557    }
558}
559
560impl<P, D, H, PE, O> XorLinkedRootedTree<P, D, H, PE, RecordEdgeChild, O>
561where
562    P: ParentComponent,
563    D: DfsPreorderComponent,
564    H: DepthComponent,
565    PE: ParentEdgeComponent,
566    O: XorBottomUpOrderComponent,
567{
568    pub fn edge_child(&self, eid: usize) -> usize {
569        self.edge_child[eid]
570    }
571    pub fn edge_children(&self) -> &[usize] {
572        &self.edge_child
573    }
574}
575
576pub struct Children<'a> {
577    dfs: &'a DfsPreorder,
578    next: usize,
579    end: usize,
580}
581
582impl Iterator for Children<'_> {
583    type Item = usize;
584    fn next(&mut self) -> Option<usize> {
585        if self.next == self.end {
586            None
587        } else {
588            let v = self.dfs.order[self.next];
589            self.next = self.dfs.subtree_end[v];
590            Some(v)
591        }
592    }
593}
594
595pub struct XorLinkedRootedTreeScanner<
596    U,
597    T = (),
598    P = NoParent,
599    D = NoDfsPreorder,
600    H = NoDepth,
601    PE = NoParentEdge,
602    EC = NoEdgeChild,
603    X = NoXorBottomUpOrder,
604    B = NoXorBottomUpOrder,
605    E = NoEIndexed,
606> where
607    U: Scan<Output = usize>,
608    T: Scan,
609{
610    n: usize,
611    root: usize,
612    _marker: ScannerMarker<U, T, P, D, H, PE, EC, X, B, E>,
613}
614
615impl<U, T> XorLinkedRootedTreeScanner<U, T>
616where
617    U: Scan<Output = usize>,
618    T: Scan,
619{
620    pub fn new(n: usize, root: usize) -> Self {
621        Self {
622            n,
623            root,
624            _marker: PhantomData,
625        }
626    }
627}
628
629impl<U, T, P, D, H, PE, EC, X, B, E> XorLinkedRootedTreeScanner<U, T, P, D, H, PE, EC, X, B, E>
630where
631    U: Scan<Output = usize>,
632    T: Scan,
633{
634    pub fn with_parent(
635        self,
636    ) -> XorLinkedRootedTreeScanner<U, T, RecordParent, D, H, PE, EC, X, B, E> {
637        XorLinkedRootedTreeScanner {
638            n: self.n,
639            root: self.root,
640            _marker: PhantomData,
641        }
642    }
643    pub fn with_dfs_preorder(
644        self,
645    ) -> XorLinkedRootedTreeScanner<
646        U,
647        T,
648        P,
649        RecordDfsPreorder,
650        H,
651        PE,
652        EC,
653        X,
654        RecordXorBottomUpOrder,
655        E,
656    > {
657        XorLinkedRootedTreeScanner {
658            n: self.n,
659            root: self.root,
660            _marker: PhantomData,
661        }
662    }
663    pub fn with_depth(
664        self,
665    ) -> XorLinkedRootedTreeScanner<U, T, P, D, RecordDepth, PE, EC, X, RecordXorBottomUpOrder, E>
666    {
667        XorLinkedRootedTreeScanner {
668            n: self.n,
669            root: self.root,
670            _marker: PhantomData,
671        }
672    }
673    pub fn with_xor_bottom_up_order(
674        self,
675    ) -> XorLinkedRootedTreeScanner<
676        U,
677        T,
678        P,
679        D,
680        H,
681        PE,
682        EC,
683        RecordXorBottomUpOrder,
684        RecordXorBottomUpOrder,
685        E,
686    > {
687        XorLinkedRootedTreeScanner {
688            n: self.n,
689            root: self.root,
690            _marker: PhantomData,
691        }
692    }
693}
694
695impl<U, T, P, D, H, X, B>
696    XorLinkedRootedTreeScanner<U, T, P, D, H, NoParentEdge, NoEdgeChild, X, B, NoEIndexed>
697where
698    U: Scan<Output = usize>,
699    T: Scan,
700{
701    pub fn with_eindexed(
702        self,
703    ) -> XorLinkedRootedTreeScanner<U, T, P, D, H, NoParentEdge, NoEdgeChild, X, B, EIndexed> {
704        XorLinkedRootedTreeScanner {
705            n: self.n,
706            root: self.root,
707            _marker: PhantomData,
708        }
709    }
710}
711
712impl<U, T, P, D, H, PE, EC, X, B> XorLinkedRootedTreeScanner<U, T, P, D, H, PE, EC, X, B, EIndexed>
713where
714    U: Scan<Output = usize>,
715    T: Scan,
716{
717    pub fn with_parent_edge(
718        self,
719    ) -> XorLinkedRootedTreeScanner<U, T, P, D, H, RecordParentEdge, EC, X, B, EIndexed> {
720        XorLinkedRootedTreeScanner {
721            n: self.n,
722            root: self.root,
723            _marker: PhantomData,
724        }
725    }
726    pub fn with_edge_child(
727        self,
728    ) -> XorLinkedRootedTreeScanner<U, T, P, D, H, PE, RecordEdgeChild, X, B, EIndexed> {
729        XorLinkedRootedTreeScanner {
730            n: self.n,
731            root: self.root,
732            _marker: PhantomData,
733        }
734    }
735}
736
737impl<U, T, P, D, H, X, B> MarkedScan
738    for XorLinkedRootedTreeScanner<U, T, P, D, H, NoParentEdge, NoEdgeChild, X, B, NoEIndexed>
739where
740    U: Scan<Output = usize>,
741    T: Scan,
742    P: ParentComponent,
743    D: DfsPreorderComponent,
744    H: DepthComponent,
745    X: BuildXorBottomUpOrder<B>,
746    B: XorBottomUpOrderBuffer,
747{
748    type Output = (
749        XorLinkedRootedTree<P, D, H, NoParentEdge, NoEdgeChild, X>,
750        Vec<<T as Scan>::Output>,
751    );
752
753    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
754        let mut acc = XorAccumulator::new(self.n);
755        let mut weights = Vec::with_capacity(self.n.saturating_sub(1));
756        for _ in 0..self.n.saturating_sub(1) {
757            let u = U::scan(iter)?;
758            let v = U::scan(iter)?;
759            acc.add_edge(u, v);
760            weights.push(T::scan(iter)?);
761        }
762        Some((
763            finish_accumulator::<P, D, H, X, B>(self.n, self.root, acc),
764            weights,
765        ))
766    }
767}
768
769impl<U, T, P, D, H, PE, EC, X, B> MarkedScan
770    for XorLinkedRootedTreeScanner<U, T, P, D, H, PE, EC, X, B, EIndexed>
771where
772    U: Scan<Output = usize>,
773    T: Scan,
774    P: ParentComponent,
775    D: DfsPreorderComponent,
776    H: DepthComponent,
777    PE: ParentEdgeComponent,
778    EC: EdgeChildComponent,
779    X: BuildXorBottomUpOrder<B>,
780    B: XorBottomUpOrderBuffer,
781{
782    type Output = (
783        XorLinkedRootedTree<P, D, H, PE, EC, X>,
784        Vec<<T as Scan>::Output>,
785    );
786
787    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
788        let mut acc = XorEIndexedAccumulator::new(self.n);
789        let mut weights = Vec::with_capacity(self.n.saturating_sub(1));
790        for eid in 0..self.n.saturating_sub(1) {
791            let u = U::scan(iter)?;
792            let v = U::scan(iter)?;
793            acc.add_edge(eid, u, v);
794            weights.push(T::scan(iter)?);
795        }
796        Some((
797            finish_eindexed_accumulator::<P, D, H, PE, EC, X, B>(self.n, self.root, acc),
798            weights,
799        ))
800    }
801}
802
803fn finish_accumulator<P, D, H, X, B>(
804    n: usize,
805    root: usize,
806    acc: XorAccumulator,
807) -> XorLinkedRootedTree<P, D, H, NoParentEdge, NoEdgeChild, X>
808where
809    P: ParentComponent,
810    D: DfsPreorderComponent,
811    H: DepthComponent,
812    X: BuildXorBottomUpOrder<B>,
813    B: XorBottomUpOrderBuffer,
814{
815    let mut xor_order = B::new(n);
816    let parent = acc.finish::<B, _>(root, &mut xor_order, |_, _| {});
817    finish_rooted_tree::<P, D, H, X, B>(root, parent, xor_order)
818}
819
820fn finish_rooted_tree<P, D, H, X, B>(
821    root: usize,
822    parent: Vec<usize>,
823    xor_order: B::Data,
824) -> XorLinkedRootedTree<P, D, H, NoParentEdge, NoEdgeChild, X>
825where
826    P: ParentComponent,
827    D: DfsPreorderComponent,
828    H: DepthComponent,
829    X: BuildXorBottomUpOrder<B>,
830    B: XorBottomUpOrderBuffer,
831{
832    let n = parent.len();
833    let order = B::as_slice(&xor_order);
834    let dfs = D::build(n, root, &parent, order);
835    let depth = H::build(n, root, &parent, order);
836    XorLinkedRootedTree {
837        n,
838        root,
839        parent: P::build(parent),
840        dfs,
841        depth,
842        parent_edge: (),
843        edge_child: (),
844        xor_order: X::build(xor_order),
845        _marker: PhantomData,
846    }
847}
848
849fn finish_eindexed_accumulator<P, D, H, PE, EC, X, B>(
850    n: usize,
851    root: usize,
852    acc: XorEIndexedAccumulator,
853) -> XorLinkedRootedTree<P, D, H, PE, EC, X>
854where
855    P: ParentComponent,
856    D: DfsPreorderComponent,
857    H: DepthComponent,
858    PE: ParentEdgeComponent,
859    EC: EdgeChildComponent,
860    X: BuildXorBottomUpOrder<B>,
861    B: XorBottomUpOrderBuffer,
862{
863    let mut xor_order = B::new(n);
864    let mut edge_child = EC::new(n.saturating_sub(1));
865    let (parent, parent_edge) =
866        acc.finish::<B, EC, _>(root, &mut xor_order, &mut edge_child, |_, _, _| {});
867    let order = B::as_slice(&xor_order);
868    let dfs = D::build(n, root, &parent, order);
869    let depth = H::build(n, root, &parent, order);
870    XorLinkedRootedTree {
871        n,
872        root,
873        parent: P::build(parent),
874        dfs,
875        depth,
876        parent_edge: PE::build(parent_edge),
877        edge_child,
878        xor_order: X::build(xor_order),
879        _marker: PhantomData,
880    }
881}
882
883struct XorAccumulator {
884    deg: Vec<isize>,
885    xor: Vec<usize>,
886}
887
888impl XorAccumulator {
889    fn new(n: usize) -> Self {
890        Self {
891            deg: vec![0; n],
892            xor: vec![0; n],
893        }
894    }
895    fn add_edge(&mut self, u: usize, v: usize) {
896        self.deg[u] += 1;
897        self.deg[v] += 1;
898        self.xor[u] ^= v;
899        self.xor[v] ^= u;
900    }
901    fn finish<O, F>(mut self, root: usize, xor_order: &mut O::Data, mut f: F) -> Vec<usize>
902    where
903        O: XorBottomUpOrderBuffer,
904        F: FnMut(usize, usize),
905    {
906        self.deg[root] = 0;
907        for i in 0..self.deg.len() {
908            let mut v = i;
909            while self.deg[v] == 1 {
910                let p = self.xor[v];
911                O::push(xor_order, v);
912                f(v, p);
913                self.deg[v] = 0;
914                self.deg[p] -= 1;
915                self.xor[p] ^= v;
916                v = p;
917            }
918        }
919        self.xor[root] = usize::MAX;
920        self.xor
921    }
922}
923
924struct XorEIndexedAccumulator {
925    deg: Vec<isize>,
926    xor: Vec<usize>,
927    edge_xor: Vec<usize>,
928}
929
930impl XorEIndexedAccumulator {
931    fn new(n: usize) -> Self {
932        Self {
933            deg: vec![0; n],
934            xor: vec![0; n],
935            edge_xor: vec![0; n],
936        }
937    }
938    fn add_edge(&mut self, eid: usize, u: usize, v: usize) {
939        self.deg[u] += 1;
940        self.deg[v] += 1;
941        self.xor[u] ^= v;
942        self.xor[v] ^= u;
943        self.edge_xor[u] ^= eid;
944        self.edge_xor[v] ^= eid;
945    }
946    fn finish<O, EC, F>(
947        mut self,
948        root: usize,
949        xor_order: &mut O::Data,
950        edge_child: &mut EC::Data,
951        mut f: F,
952    ) -> (Vec<usize>, Vec<usize>)
953    where
954        O: XorBottomUpOrderBuffer,
955        EC: EdgeChildComponent,
956        F: FnMut(usize, usize, usize),
957    {
958        self.deg[root] = 0;
959        for i in 0..self.deg.len() {
960            let mut v = i;
961            while self.deg[v] == 1 {
962                let p = self.xor[v];
963                let e = self.edge_xor[v];
964                O::push(xor_order, v);
965                EC::set(edge_child, e, v);
966                f(v, p, e);
967                self.deg[v] = 0;
968                self.deg[p] -= 1;
969                self.xor[p] ^= v;
970                self.edge_xor[p] ^= e;
971                v = p;
972            }
973        }
974        self.xor[root] = usize::MAX;
975        self.edge_xor[root] = usize::MAX;
976        (self.xor, self.edge_xor)
977    }
978}
979
980#[cfg(test)]
981mod tests {
982    use super::*;
983    use crate::{
984        graph::{Graph, UndirectedSparseGraph},
985        scan,
986        tools::{Scanner, Xorshift},
987        tree::MixedTree,
988    };
989
990    fn expected_parent_depth(
991        graph: &UndirectedSparseGraph,
992        root: usize,
993    ) -> (Vec<usize>, Vec<usize>) {
994        let n = graph.vertices_size();
995        let mut parent = vec![usize::MAX; n];
996        let mut depth = vec![0usize; n];
997        let mut stack = vec![root];
998        while let Some(u) = stack.pop() {
999            for a in graph.neighbors(u) {
1000                if a.to != parent[u] {
1001                    parent[a.to] = u;
1002                    depth[a.to] = depth[u] + 1;
1003                    stack.push(a.to);
1004                }
1005            }
1006        }
1007        (parent, depth)
1008    }
1009
1010    fn assert_rooted_tree<O>(
1011        graph: &UndirectedSparseGraph,
1012        root: usize,
1013        tree: &XorLinkedRootedTree<
1014            RecordParent,
1015            RecordDfsPreorder,
1016            RecordDepth,
1017            NoParentEdge,
1018            NoEdgeChild,
1019            O,
1020        >,
1021    ) where
1022        O: XorBottomUpOrderComponent,
1023    {
1024        let n = graph.vertices_size();
1025        let (parent, depth) = expected_parent_depth(graph, root);
1026        assert_eq!(tree.vertices_size(), n);
1027        assert_eq!(tree.edges_size(), n.saturating_sub(1));
1028        assert_eq!(tree.root(), root);
1029        assert_eq!(tree.parents(), parent);
1030        assert_eq!(tree.depths(), depth);
1031
1032        let mut seen = vec![false; n];
1033        for (i, &v) in tree.dfs_order().iter().enumerate() {
1034            assert!(!seen[v]);
1035            seen[v] = true;
1036            assert_eq!(tree.dfs_index(v), i);
1037        }
1038        assert!(seen.into_iter().all(|x| x));
1039
1040        let mut children = vec![vec![]; n];
1041        for v in 0..n {
1042            if v != root {
1043                children[parent[v]].push(v);
1044            }
1045        }
1046        for (v, expected_children) in children.iter_mut().enumerate() {
1047            let mut actual: Vec<_> = tree.children(v).collect();
1048            actual.sort_unstable();
1049            expected_children.sort_unstable();
1050            assert_eq!(actual, *expected_children);
1051            assert_eq!(
1052                tree.subtree_size(v),
1053                expected_children
1054                    .iter()
1055                    .map(|&u| tree.subtree_size(u))
1056                    .sum::<usize>()
1057                    + 1
1058            );
1059            let range = tree.subtree_range(v);
1060            for &u in &tree.dfs_order()[range.clone()] {
1061                let mut x = u;
1062                while x != v && x != usize::MAX {
1063                    x = parent[x];
1064                }
1065                assert_eq!(x, v);
1066            }
1067            assert_eq!(range.len(), tree.subtree_size(v));
1068        }
1069    }
1070
1071    #[test]
1072    fn xor_linked_tree_rooted_properties() {
1073        let mut rng = Xorshift::default();
1074        for n in 1..=6 {
1075            for parents in crate::tools::testutil::exhaustive_sequences(0..n, n - 1..=n - 1) {
1076                if parents.iter().enumerate().all(|(i, &p)| p <= i) {
1077                    let edges = parents
1078                        .iter()
1079                        .enumerate()
1080                        .map(|(i, &p)| (p, i + 1))
1081                        .collect();
1082                    let graph = UndirectedSparseGraph::from_edges(n, edges);
1083                    let tree = XorLinkedRootedTree::builder(n)
1084                        .with_parent()
1085                        .with_dfs_preorder()
1086                        .with_depth()
1087                        .build_from_ordered_parents(parents);
1088                    assert_rooted_tree(&graph, 0, &tree);
1089                }
1090            }
1091        }
1092        for n in 1..=200 {
1093            for kind in 0..3 {
1094                let graph = rng.random(MixedTree(n));
1095                let root = rng.random(0..n);
1096                let tree = XorLinkedRootedTree::builder(n)
1097                    .with_parent()
1098                    .with_dfs_preorder()
1099                    .with_depth()
1100                    .build(root, graph.edges.iter().copied());
1101                assert_rooted_tree(&graph, root, &tree);
1102                let parents: Vec<_> = (1..n)
1103                    .map(|v| match kind {
1104                        0 => 0,
1105                        1 => v - 1,
1106                        _ => rng.random(0..v),
1107                    })
1108                    .collect();
1109                let edges = parents
1110                    .iter()
1111                    .enumerate()
1112                    .map(|(i, &p)| (p, i + 1))
1113                    .collect();
1114                let graph = UndirectedSparseGraph::from_edges(n, edges);
1115                let tree = XorLinkedRootedTree::builder(n)
1116                    .with_parent()
1117                    .with_dfs_preorder()
1118                    .with_depth()
1119                    .build_from_ordered_parents(parents);
1120                assert_rooted_tree(&graph, 0, &tree);
1121            }
1122        }
1123    }
1124
1125    #[test]
1126    fn xor_linked_tree_eindexed_properties() {
1127        let mut rng = Xorshift::default();
1128        for n in 1..=200 {
1129            for _ in 0..3 {
1130                let graph = rng.random(MixedTree(n));
1131                let root = rng.random(0..n);
1132                let tree = XorLinkedRootedTree::builder(n)
1133                    .with_parent()
1134                    .with_eindexed()
1135                    .with_parent_edge()
1136                    .with_edge_child()
1137                    .build(root, graph.edges.iter().copied());
1138                assert_eq!(tree.parent(root), usize::MAX);
1139                assert_eq!(tree.parent_edge(root), usize::MAX);
1140                for (eid, &(u, v)) in graph.edges.iter().enumerate() {
1141                    let child = if tree.parent(u) == v {
1142                        u
1143                    } else {
1144                        assert_eq!(tree.parent(v), u);
1145                        v
1146                    };
1147                    assert_eq!(tree.parent_edge(child), eid);
1148                    assert_eq!(tree.edge_child(eid), child);
1149                }
1150            }
1151        }
1152    }
1153
1154    #[test]
1155    fn xor_linked_tree_xor_bottom_up_order() {
1156        let mut rng = Xorshift::default();
1157        for n in 1..=200 {
1158            for _ in 0..3 {
1159                let graph = rng.random(MixedTree(n));
1160                let root = rng.random(0..n);
1161                let tree = XorLinkedRootedTree::builder(n)
1162                    .with_parent()
1163                    .with_xor_bottom_up_order()
1164                    .build(root, graph.edges.iter().copied());
1165                let bottom_up = tree.xor_bottom_up_order();
1166                assert_eq!(bottom_up.len(), n - 1);
1167                assert!(!bottom_up.contains(&root));
1168
1169                let mut bottom_up_index = vec![usize::MAX; n];
1170                for (i, &v) in bottom_up.iter().enumerate() {
1171                    assert_eq!(bottom_up_index[v], usize::MAX);
1172                    bottom_up_index[v] = i;
1173                }
1174                for v in 0..n {
1175                    if v != root && tree.parent(v) != root {
1176                        assert!(bottom_up_index[v] < bottom_up_index[tree.parent(v)]);
1177                    }
1178                }
1179
1180                let top_down: Vec<_> = tree.xor_top_down_order().collect();
1181                assert_eq!(
1182                    top_down,
1183                    bottom_up.iter().rev().copied().collect::<Vec<_>>()
1184                );
1185                let mut depth = vec![0usize; n];
1186                for v in top_down {
1187                    depth[v] = depth[tree.parent(v)] + 1;
1188                }
1189                let (_, expected_depth) = expected_parent_depth(&graph, root);
1190                assert_eq!(depth, expected_depth);
1191            }
1192        }
1193    }
1194
1195    #[test]
1196    fn xor_linked_tree_visitor() {
1197        let mut rng = Xorshift::default();
1198        for _ in 0..1000 {
1199            let n = rng.random(1..=100);
1200            let graph = rng.random(MixedTree(n));
1201            let root = rng.random(0..n);
1202            let (parent, _) = expected_parent_depth(&graph, root);
1203            let mut seen = Vec::new();
1204            XorLinkedRootedTree::builder(n)
1205                .run(root, graph.edges.iter().copied(), |u, p| seen.push((u, p)));
1206            seen.sort();
1207            assert_eq!(
1208                seen,
1209                (0..n)
1210                    .filter(|&u| u != root)
1211                    .map(|u| (u, parent[u]))
1212                    .collect::<Vec<_>>()
1213            );
1214            let mut seen_e = Vec::new();
1215            XorLinkedRootedTree::builder(n).with_eindexed().run(
1216                root,
1217                graph.edges.iter().copied(),
1218                |u, p, e| seen_e.push((u, p, e)),
1219            );
1220            seen_e.sort_by_key(|&(_, _, e)| e);
1221            assert_eq!(seen_e.len(), graph.edges_size());
1222            for (eid, &(u, p, e)) in seen_e.iter().enumerate() {
1223                assert_eq!(e, eid);
1224                assert_eq!(parent[u], p);
1225                assert!(graph.edges[e] == (u, p) || graph.edges[e] == (p, u));
1226            }
1227        }
1228    }
1229
1230    #[derive(Debug, PartialEq, Eq)]
1231    struct NonCloneWeight(usize);
1232
1233    impl Scan for NonCloneWeight {
1234        type Output = NonCloneWeight;
1235        fn scan<I: ScanSource>(iter: &mut I) -> Option<NonCloneWeight> {
1236            Some(NonCloneWeight(usize::scan(iter)?))
1237        }
1238    }
1239
1240    #[test]
1241    fn xor_linked_tree_scanner() {
1242        let mut rng = Xorshift::default();
1243        for _ in 0..1000 {
1244            let n = rng.random(1..=100);
1245            let graph = rng.random(MixedTree(n));
1246            let root = rng.random(0..n);
1247            let (parent, _) = expected_parent_depth(&graph, root);
1248            let expected: Vec<usize> = rng.random_iter(..).take(graph.edges_size()).collect();
1249            let text: String = graph
1250                .edges
1251                .iter()
1252                .zip(&expected)
1253                .map(|(&(u, v), w)| format!("{u} {v} {w}\n"))
1254                .collect();
1255            let mut scanner = Scanner::new(&text);
1256            scan!(scanner, (tree, weights): @XorLinkedRootedTreeScanner::<usize, NonCloneWeight>::new(n, root).with_parent().with_eindexed().with_parent_edge());
1257            assert_eq!(
1258                weights,
1259                expected.into_iter().map(NonCloneWeight).collect::<Vec<_>>()
1260            );
1261            for (u, &p) in parent.iter().enumerate() {
1262                assert_eq!(tree.parent(u), p);
1263            }
1264            assert_eq!(tree.parent_edge(root), usize::MAX);
1265            for (eid, &(u, v)) in graph.edges.iter().enumerate() {
1266                let child = if parent[u] == v { u } else { v };
1267                assert_eq!(tree.parent_edge(child), eid);
1268            }
1269        }
1270    }
1271}
competitive/tree/xor_linked_tree.rs

competitive/tree/
xor_linked_tree.rs
