library_checker/tree/
dynamic_tree_vertex_add_subtree_sum.rs

1use competitive::prelude::*;
2use competitive::tree::{
3    LinkCutTree, LinkCutTreeSpec, LinkCutTreeSubtreeFold, TopTree, TopTreeSpec,
4};
5
6struct SubtreeSum;
7
8struct SubtreeSumData {
9    value: u64,
10    virtual_sum: u64,
11    sum: u64,
12}
13
14impl LinkCutTreeSpec for SubtreeSum {
15    type Value = u64;
16    type Data = SubtreeSumData;
17
18    const ROOT_TO_NODE_TOP_DOWN: bool = false;
19
20    fn new(value: Self::Value) -> Self::Data {
21        SubtreeSumData {
22            value,
23            virtual_sum: 0,
24            sum: value,
25        }
26    }
27
28    fn value(data: &Self::Data) -> &Self::Value {
29        &data.value
30    }
31
32    fn value_mut(data: &mut Self::Data) -> &mut Self::Value {
33        &mut data.value
34    }
35
36    fn bottom_up(data: &mut Self::Data, children: [Option<&Self::Data>; 2]) {
37        data.sum = data.value
38            + data.virtual_sum
39            + children
40                .into_iter()
41                .flatten()
42                .map(|child| child.sum)
43                .sum::<u64>();
44    }
45
46    fn reverse(_data: &mut Self::Data) {}
47
48    fn attach_virtual(parent: &mut Self::Data, child: &mut Self::Data) {
49        parent.virtual_sum += child.sum;
50    }
51
52    fn detach_virtual(parent: &mut Self::Data, child: &mut Self::Data) {
53        parent.virtual_sum -= child.sum;
54    }
55}
56
57impl LinkCutTreeSubtreeFold for SubtreeSum {
58    type Subtree = u64;
59
60    fn fold_subtree(data: &Self::Data) -> Self::Subtree {
61        data.sum
62    }
63}
64
65impl TopTreeSpec for SubtreeSum {
66    type Info = u64;
67    type Point = (u64, u64);
68    type Path = (u64, u64, u64);
69
70    fn vertex(info: &Self::Info) -> Self::Path {
71        (*info, 1, *info)
72    }
73
74    fn add_vertex(point: &Self::Point, info: &Self::Info) -> Self::Path {
75        (point.0 + *info, point.1 + 1, *info)
76    }
77
78    fn add_edge(path: &Self::Path) -> Self::Point {
79        (path.0, path.1)
80    }
81
82    fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point {
83        (left.0 + right.0, left.1 + right.1)
84    }
85
86    fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path {
87        (left.0 + right.0, left.1 + right.1, left.2 + right.2)
88    }
89
90    fn reverse(_path: &mut Self::Path) {}
91}
92
93competitive::define_enum_scan! {
94    enum Query: usize {
95        0 => Relink { u: usize, v: usize, w: usize, x: usize }
96        1 => Add { p: usize, x: u64 }
97        2 => Sum { v: usize, p: usize }
98    }
99}
100
101#[verify::library_checker("dynamic_tree_vertex_add_subtree_sum")]
102pub fn dynamic_tree_vertex_add_subtree_sum(reader: impl Read, writer: impl Write) {
103    prepare_io!(reader, writer);
104    sc!(n, q, a: [u64; n], edges: [(usize, usize); n - 1]);
105    let mut tree = LinkCutTree::<SubtreeSum>::from_edges(a, &edges);
106    for _ in 0..q {
107        sc!(query: Query);
108        match query {
109            Query::Relink { u, v, w, x } => {
110                tree.cut(u, v);
111                tree.link(w, x);
112            }
113            Query::Add { p, x } => tree.modify(p, |value| *value + x),
114            Query::Sum { v, p } => {
115                pp!(tree.fold_subtree(v, p));
116            }
117        }
118    }
119}
120
121#[verify::library_checker("dynamic_tree_vertex_add_subtree_sum")]
122pub fn dynamic_tree_vertex_add_subtree_sum_top_tree(reader: impl Read, writer: impl Write) {
123    prepare_io!(reader, writer);
124    sc!(n, q, a: [u64; n], edges: [(usize, usize); n - 1]);
125    let mut tree = TopTree::<SubtreeSum>::from_edges(a, &edges);
126    for _ in 0..q {
127        sc!(query: Query);
128        match query {
129            Query::Relink { u, v, w, x } => {
130                tree.cut(u, v);
131                tree.link(w, x);
132            }
133            Query::Add { p, x } => tree.modify(p, |value| *value + x),
134            Query::Sum { v, p } => {
135                pp!(tree.fold_subtree(v, p).0);
136            }
137        }
138    }
139}
library_checker/tree/dynamic_tree_vertex_add_subtree_sum.rs

library_checker/tree/
dynamic_tree_vertex_add_subtree_sum.rs
