library_checker/datastructure/
vertex_add_subtree_sum.rs

1use competitive::prelude::*;
2#[doc(no_inline)]
3pub use competitive::{
4    algebra::AdditiveOperation, data_structure::SegmentTree, graph::UndirectedSparseGraph,
5    tree::EulerTourForVertex,
6};
7
8#[verify::library_checker("vertex_add_subtree_sum")]
9pub fn vertex_add_subtree_sum(reader: impl Read, mut writer: impl Write) {
10    let s = read_all_unchecked(reader);
11    let mut scanner = Scanner::new(&s);
12    scan!(scanner, n, q, a: [u64; n], p: [usize]);
13    let edges = p.take(n - 1).enumerate().map(|(i, p)| (i + 1, p)).collect();
14    let graph = UndirectedSparseGraph::from_edges(n, edges);
15    let mut et = EulerTourForVertex::new(&graph);
16    et.subtree_vertex_tour(0, n);
17    let mut b = vec![0; n];
18    for i in 0..n {
19        b[et.vidx[i].0] = a[i];
20    }
21    let mut seg = SegmentTree::<AdditiveOperation<_>>::from_vec(b);
22    for _ in 0..q {
23        match scanner.scan::<usize>() {
24            0 => {
25                scan!(scanner, u, x: u64);
26                et.subtree_update(u, x, |k, x| seg.update(k, x));
27            }
28            1 => {
29                scan!(scanner, u);
30                writeln!(writer, "{}", et.subtree_query(u, |l, r| seg.fold(l..r))).ok();
31            }
32            _ => panic!("unknown query"),
33        }
34    }
35}