Skip to main content

library_checker/tree/
vertex_get_range_contour_add_on_tree.rs

1use 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 { v: usize, l: usize, r: usize, x: i64 }
9        1 => Get { v: usize }
10    }
11}
12
13#[verify::library_checker("vertex_get_range_contour_add_on_tree")]
14pub fn vertex_get_range_contour_add_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 bits: Vec<BinaryIndexedTree<AdditiveOperation<_>>> = cq
19        .component_sizes()
20        .map(|n| BinaryIndexedTree::new(n + 1))
21        .collect();
22
23    for _ in 0..q {
24        sc!(query: Query);
25        match query {
26            Query::Add { v, l, r, x } => {
27                cq.for_each_contour_range(v, l, r, |c, start, end| {
28                    bits[c].update(start, x);
29                    bits[c].update(end, -x);
30                });
31                if l == 0 && 0 < r {
32                    a[v] += x;
33                }
34            }
35            Query::Get { v } => {
36                let mut ans = a[v];
37                cq.for_each_index(v, |c, i| ans += bits[c].accumulate(i));
38                pp!(ans);
39            }
40        }
41    }
42}