library_checker/tree/
vertex_set_path_composite.rs1use competitive::prelude::*;
2use competitive::{
3 algebra::LinearOperation, graph::TreeGraphScanner, num::mint_basic::MInt998244353 as M,
4};
5
6competitive::define_enum_scan! {
7 enum Query: usize {
8 0 => Set { p: usize, cd: (M, M) }
9 1 => Apply { u: usize, v: usize, x: M }
10 }
11}
12
13#[verify::library_checker("vertex_set_path_composite")]
14pub fn vertex_set_path_composite(reader: impl Read, writer: impl Write) {
15 prepare_io!(reader, writer);
16 sc!(n, q, ab: [(M, M); n], (graph, _): @TreeGraphScanner::<usize, ()>::new(n));
17 let hld = graph.hld(0);
18 let mut fold = hld.build_fold::<LinearOperation<_>>(&ab);
19 for _ in 0..q {
20 sc!(query: Query);
21 match query {
22 Query::Set { p, cd } => {
23 fold.set(p, cd);
24 }
25 Query::Apply { u, v, x } => {
26 let (a, b) = fold.fold_vertices(u, v);
27 pp!(a * x + b);
28 }
29 }
30 }
31}