library_checker/tree/
vertex_add_path_sum.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation,
4    data_structure::BinaryIndexedTree,
5    tree::{LowestCommonAncestor, XorLinkedRootedTreeScanner},
6};
7
8competitive::define_enum_scan! {
9    enum Query: usize {
10        0 => Add { p: usize, x: i64 }
11        1 => Sum { u: usize, v: usize }
12    }
13}
14
15#[verify::library_checker("vertex_add_path_sum")]
16pub fn vertex_add_path_sum(reader: impl Read, writer: impl Write) {
17    prepare_io!(reader, writer);
18    sc!(n, q, mut a: [i64; n],
19        (tree, _): @XorLinkedRootedTreeScanner::<usize, ()>::new(n, 0)
20            .with_parent().with_dfs_preorder());
21    let lca = LowestCommonAncestor::from_dfs_preorder(tree.parents(), tree.dfs_order());
22    let mut values = vec![0; n + 1];
23    for (u, &x) in a.iter().enumerate() {
24        let range = tree.subtree_range(u);
25        values[range.start] += x;
26        values[range.end] -= x;
27    }
28    let mut bit = BinaryIndexedTree::<AdditiveOperation<_>>::from_slice(&values);
29    for _ in 0..q {
30        sc!(query: Query);
31        match query {
32            Query::Add { p, x } => {
33                a[p] += x;
34                let range = tree.subtree_range(p);
35                bit.update(range.start, x);
36                bit.update(range.end, -x);
37            }
38            Query::Sum { u, v } => {
39                let p = lca.lca(u, v);
40                pp!(
41                    a[p] + bit.accumulate(tree.dfs_index(u)) + bit.accumulate(tree.dfs_index(v))
42                        - 2 * bit.accumulate(tree.dfs_index(p))
43                );
44            }
45        }
46    }
47}
library_checker/tree/vertex_add_path_sum.rs

library_checker/tree/
vertex_add_path_sum.rs
