library_checker/tree/
vertex_add_subtree_sum.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation,
4    data_structure::{DaryPrefixSumTreeU64, SegmentTree},
5    graph::UndirectedSparseGraph,
6    tree::XorLinkedRootedTree,
7};
8
9competitive::define_enum_scan! {
10    enum Query: usize {
11        0 => Add { u: usize, x: u64 }
12        1 => Sum { u: usize }
13    }
14}
15
16#[verify::library_checker("vertex_add_subtree_sum")]
17pub fn vertex_add_subtree_sum(reader: impl Read, writer: impl Write) {
18    prepare_io!(reader, writer);
19    sc!(n, q, a: [u64; n], p: [usize; iter n - 1]);
20    let tree = XorLinkedRootedTree::builder(n)
21        .with_dfs_preorder()
22        .build_from_ordered_parents(p);
23    let b: Vec<_> = tree.dfs_order().iter().map(|&v| a[v]).collect();
24    let mut seg = DaryPrefixSumTreeU64::from_slice(&b);
25    for _ in 0..q {
26        sc!(query: Query);
27        match query {
28            Query::Add { u, x } => seg.update(tree.dfs_index(u), x),
29            Query::Sum { u } => {
30                let range = tree.subtree_range(u);
31                pp!(seg.fold(range.start, range.end));
32            }
33        }
34    }
35}
36
37#[verify::library_checker("vertex_add_subtree_sum")]
38pub fn vertex_add_subtree_sum_hld(reader: impl Read, writer: impl Write) {
39    prepare_io!(reader, writer);
40    sc!(n, q, a: [u64; n], p: [usize; iter n - 1]);
41    let edges = p.enumerate().map(|(i, p)| (i + 1, p)).collect();
42    let tree = UndirectedSparseGraph::from_edges(n, edges);
43    let hld = tree.hld(0);
44    let mut b = vec![0; n];
45    for (v, x) in a.into_iter().enumerate() {
46        b[hld.index(v)] = x;
47    }
48    let mut seg = SegmentTree::<AdditiveOperation<_>>::from_vec(b);
49    for _ in 0..q {
50        sc!(query: Query);
51        match query {
52            Query::Add { u, x } => seg.update(hld.index(u), x),
53            Query::Sum { u } => {
54                pp!(seg.fold(hld.subtree_range(u)));
55            }
56        }
57    }
58}
library_checker/tree/vertex_add_subtree_sum.rs

library_checker/tree/
vertex_add_subtree_sum.rs
