library_checker/tree/
vertex_set_path_composite.rs1use competitive::prelude::*;
2#[doc(no_inline)]
3pub use competitive::{
4 algebra::{LinearOperation, ReverseOperation},
5 data_structure::SegmentTree,
6 graph::TreeGraphScanner,
7 num::{MInt, mint_basic::MInt998244353},
8 tree::HeavyLightDecomposition,
9};
10
11competitive::define_enum_scan! {
12 enum Query: usize {
13 0 => Set { p: usize, cd: (MInt998244353, MInt998244353) }
14 1 => Apply { u: usize, v: usize, x: MInt998244353 }
15 }
16}
17
18#[verify::library_checker("vertex_set_path_composite")]
19pub fn vertex_set_path_composite(reader: impl Read, mut writer: impl Write) {
20 let s = read_all_unchecked(reader);
21 let mut scanner = Scanner::new(&s);
22 scan!(scanner, n, q, ab: [(MInt998244353, MInt998244353); n], (mut graph, _): @TreeGraphScanner::<usize, ()>::new(n));
23 let hld = HeavyLightDecomposition::new(0, &mut graph);
24 let mut nab = vec![(MInt998244353::default(), MInt998244353::default()); n];
25 for i in 0..n {
26 nab[hld.vidx[i]] = ab[i];
27 }
28 let mut seg1 = SegmentTree::<LinearOperation<_>>::from_vec(nab.clone());
29 let mut seg2 = SegmentTree::<ReverseOperation<LinearOperation<_>>>::from_vec(nab);
30 for _ in 0..q {
31 scan!(scanner, query: Query);
32 match query {
33 Query::Set { p, cd } => {
34 seg1.set(hld.vidx[p], cd);
35 seg2.set(hld.vidx[p], cd);
36 }
37 Query::Apply { u, v, x } => {
38 let (a, b) = hld.query_noncom::<LinearOperation<_>, _, _>(
39 u,
40 v,
41 false,
42 |l, r| seg1.fold(l..r),
43 |l, r| seg2.fold(l..r),
44 );
45 writeln!(writer, "{}", a * x + b).ok();
46 }
47 }
48 }
49}