library_checker/tree/
dynamic_tree_subtree_add_subtree_sum.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation,
4    tree::{
5        LinkCutTree, LinkCutTreeSpec, LinkCutTreeSubtreeFold, LinkCutTreeSubtreeUpdate, TopTree,
6        TopTreeAction, TopTreeSpec,
7    },
8};
9
10struct SubtreeSum;
11
12struct SubtreeSumData {
13    value: u64,
14    virtual_sum: u64,
15    virtual_size: u64,
16    sum: u64,
17    size: u64,
18    lazy: u64,
19    virtual_lazy: u64,
20    path_parent_lazy: u64,
21}
22
23impl SubtreeSum {
24    fn apply(data: &mut SubtreeSumData, action: u64) {
25        if action == 0 {
26            return;
27        }
28        data.value += action;
29        data.virtual_sum += data.virtual_size * action;
30        data.sum += data.size * action;
31        data.lazy += action;
32        data.virtual_lazy += action;
33    }
34}
35
36impl LinkCutTreeSpec for SubtreeSum {
37    type Value = u64;
38    type Data = SubtreeSumData;
39
40    const ROOT_TO_NODE_TOP_DOWN: bool = false;
41
42    fn new(value: Self::Value) -> Self::Data {
43        SubtreeSumData {
44            value,
45            virtual_sum: 0,
46            virtual_size: 0,
47            sum: value,
48            size: 1,
49            lazy: 0,
50            virtual_lazy: 0,
51            path_parent_lazy: 0,
52        }
53    }
54
55    fn value(data: &Self::Data) -> &Self::Value {
56        &data.value
57    }
58
59    fn value_mut(data: &mut Self::Data) -> &mut Self::Value {
60        &mut data.value
61    }
62
63    fn top_down(data: &mut Self::Data, children: [Option<&mut Self::Data>; 2]) {
64        let action = std::mem::replace(&mut data.lazy, 0);
65        for child in children.into_iter().flatten() {
66            Self::apply(child, action);
67        }
68    }
69
70    fn bottom_up(data: &mut Self::Data, children: [Option<&Self::Data>; 2]) {
71        data.sum = data.value
72            + data.virtual_sum
73            + children
74                .into_iter()
75                .flatten()
76                .map(|child| child.sum)
77                .sum::<u64>();
78        data.size = 1
79            + data.virtual_size
80            + children
81                .into_iter()
82                .flatten()
83                .map(|child| child.size)
84                .sum::<u64>();
85    }
86
87    fn reverse(_data: &mut Self::Data) {}
88
89    fn attach_virtual(parent: &mut Self::Data, child: &mut Self::Data) {
90        child.path_parent_lazy = parent.virtual_lazy;
91        parent.virtual_sum += child.sum;
92        parent.virtual_size += child.size;
93    }
94
95    fn detach_virtual(parent: &mut Self::Data, child: &mut Self::Data) {
96        Self::apply(child, parent.virtual_lazy - child.path_parent_lazy);
97        parent.virtual_sum -= child.sum;
98        parent.virtual_size -= child.size;
99    }
100
101    fn transfer_path_parent(old_root: &mut Self::Data, new_root: &mut Self::Data) {
102        new_root.path_parent_lazy = std::mem::replace(&mut old_root.path_parent_lazy, 0);
103    }
104}
105
106impl LinkCutTreeSubtreeFold for SubtreeSum {
107    type Subtree = u64;
108
109    fn fold_subtree(data: &Self::Data) -> Self::Subtree {
110        data.sum
111    }
112}
113
114impl LinkCutTreeSubtreeUpdate for SubtreeSum {
115    type SubtreeAction = u64;
116
117    fn update_subtree(data: &mut Self::Data, action: &Self::SubtreeAction) {
118        Self::apply(data, *action);
119    }
120}
121
122struct TopTreeSubtreeSum;
123
124impl TopTreeSpec for TopTreeSubtreeSum {
125    type Info = u64;
126    type Point = (u64, u64);
127    type Path = (u64, u64, u64);
128
129    fn vertex(info: &Self::Info) -> Self::Path {
130        (*info, 1, 1)
131    }
132
133    fn add_vertex(point: &Self::Point, info: &Self::Info) -> Self::Path {
134        (point.0 + *info, point.1 + 1, 1)
135    }
136
137    fn add_edge(path: &Self::Path) -> Self::Point {
138        (path.0, path.1)
139    }
140
141    fn rake(left: &Self::Point, right: &Self::Point) -> Self::Point {
142        (left.0 + right.0, left.1 + right.1)
143    }
144
145    fn compress(left: &Self::Path, right: &Self::Path) -> Self::Path {
146        (left.0 + right.0, left.1 + right.1, left.2 + right.2)
147    }
148
149    fn reverse(_path: &mut Self::Path) {}
150}
151
152struct AddAction;
153
154impl TopTreeAction<TopTreeSubtreeSum> for AddAction {
155    type Action = u64;
156    type ActionMonoid = AdditiveOperation<u64>;
157
158    fn act_info(info: &mut u64, action: &Self::Action) {
159        *info += *action;
160    }
161
162    fn act_point(point: &mut (u64, u64), action: &Self::Action) {
163        point.0 += point.1 * *action;
164    }
165
166    fn act_path(path: &mut (u64, u64, u64), action: &Self::Action) {
167        path.0 += path.2 * *action;
168    }
169
170    fn act_path_light(path: &mut (u64, u64, u64), action: &Self::Action) {
171        path.0 += (path.1 - path.2) * *action;
172    }
173}
174
175competitive::define_enum_scan! {
176    enum Query: usize {
177        0 => Relink { u: usize, v: usize, w: usize, x: usize }
178        1 => Add { v: usize, p: usize, x: u64 }
179        2 => Sum { v: usize, p: usize }
180    }
181}
182
183#[verify::library_checker("dynamic_tree_subtree_add_subtree_sum")]
184pub fn dynamic_tree_subtree_add_subtree_sum(reader: impl Read, writer: impl Write) {
185    prepare_io!(reader, writer);
186    sc!(n, q, a: [u64; n], edges: [(usize, usize); n - 1]);
187    let mut tree = LinkCutTree::<SubtreeSum>::from_edges(a, &edges);
188    for _ in 0..q {
189        sc!(query: Query);
190        match query {
191            Query::Relink { u, v, w, x } => {
192                tree.cut(u, v);
193                tree.link(w, x);
194            }
195            Query::Add { v, p, x } => tree.update_subtree(v, p, &x),
196            Query::Sum { v, p } => {
197                pp!(tree.fold_subtree(v, p));
198            }
199        }
200    }
201}
202
203#[verify::library_checker("dynamic_tree_subtree_add_subtree_sum")]
204pub fn dynamic_tree_subtree_add_subtree_sum_top_tree(reader: impl Read, writer: impl Write) {
205    prepare_io!(reader, writer);
206    sc!(n, q, a: [u64; n], edges: [(usize, usize); n - 1]);
207    let mut tree = TopTree::<TopTreeSubtreeSum, AddAction>::from_edges(a, &edges);
208    for _ in 0..q {
209        sc!(query: Query);
210        match query {
211            Query::Relink { u, v, w, x } => {
212                tree.cut(u, v);
213                tree.link(w, x);
214            }
215            Query::Add { v, p, x } => tree.update_subtree(v, p, &x),
216            Query::Sum { v, p } => {
217                pp!(tree.fold_subtree(v, p).0);
218            }
219        }
220    }
221}
library_checker/tree/dynamic_tree_subtree_add_subtree_sum.rs

library_checker/tree/
dynamic_tree_subtree_add_subtree_sum.rs
