Skip to main content

library_checker/tree/
vertex_set_path_composite.rs

1use 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}