library_checker/datastructure/
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
11#[verify::library_checker("vertex_set_path_composite")]
12pub fn vertex_set_path_composite(reader: impl Read, mut writer: impl Write) {
13    let s = read_all_unchecked(reader);
14    let mut scanner = Scanner::new(&s);
15    scan!(scanner, n, q, ab: [(MInt998244353, MInt998244353); n], (mut graph, _): @TreeGraphScanner::<usize, ()>::new(n));
16    let hld = HeavyLightDecomposition::new(0, &mut graph);
17    let mut nab = vec![(MInt998244353::default(), MInt998244353::default()); n];
18    for i in 0..n {
19        nab[hld.vidx[i]] = ab[i];
20    }
21    let mut seg1 = SegmentTree::<LinearOperation<_>>::from_vec(nab.clone());
22    let mut seg2 = SegmentTree::<ReverseOperation<LinearOperation<_>>>::from_vec(nab);
23    for _ in 0..q {
24        match scanner.scan::<usize>() {
25            0 => {
26                scan!(scanner, p, cd: (MInt998244353, MInt998244353));
27                seg1.set(hld.vidx[p], cd);
28                seg2.set(hld.vidx[p], cd);
29            }
30            1 => {
31                scan!(scanner, u, v, x: MInt998244353);
32                let (a, b) = hld.query_noncom::<LinearOperation<_>, _, _>(
33                    u,
34                    v,
35                    false,
36                    |l, r| seg1.fold(l..r),
37                    |l, r| seg2.fold(l..r),
38                );
39                writeln!(writer, "{}", a * x + b).ok();
40            }
41            _ => panic!("unknown query"),
42        }
43    }
44}