library_checker/tree/
vertex_add_range_contour_sum_on_tree.rs1use competitive::prelude::*;
2use competitive::{
3 algebra::AdditiveOperation, data_structure::BinaryIndexedTree, graph::TreeGraphScanner,
4};
5
6competitive::define_enum_scan! {
7 enum Query: usize {
8 0 => Add { p: usize, x: i64 }
9 1 => Sum { v: usize, l: usize, r: usize }
10 }
11}
12
13#[verify::library_checker("vertex_add_range_contour_sum_on_tree")]
14pub fn vertex_add_range_contour_sum_on_tree(reader: impl Read, writer: impl Write) {
15 prepare_io!(reader, writer);
16 sc!(n, q, mut a: [i64; n], (graph, _): @TreeGraphScanner::<usize, ()>::new(n));
17 let cq = graph.contour_query_range();
18 let mut raw: Vec<_> = cq.component_sizes().map(|n| vec![0; n]).collect();
19 for (v, &x) in a.iter().enumerate() {
20 cq.for_each_index(v, |c, i| raw[c][i] += x);
21 }
22 let mut bits: Vec<BinaryIndexedTree<AdditiveOperation<_>>> = raw
23 .into_iter()
24 .map(|values| BinaryIndexedTree::from_slice(&values))
25 .collect();
26 for _ in 0..q {
27 sc!(query: Query);
28 match query {
29 Query::Add { p, x } => {
30 a[p] += x;
31 cq.for_each_index(p, |c, i| bits[c].update(i, x));
32 }
33 Query::Sum { v, l, r } => {
34 let mut ans = if l == 0 && 0 < r { a[v] } else { 0 };
35 cq.for_each_contour_range(v, l, r, |c, start, end| {
36 ans += bits[c].fold_abelian(start, end);
37 });
38 pp!(ans);
39 }
40 }
41 }
42}