Skip to main content

library_checker/data_structure/
range_affine_range_sum_large_array.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::RangeSumRangeLinear,
4    algorithm::SliceSortExt,
5    data_structure::{LazySegmentTree, StaticSearch},
6    num::{Zero, mint_basic::MInt998244353 as M},
7};
8
9competitive::define_enum_scan! {
10    #[derive(Clone, Copy)]
11    enum Query: u8 {
12        0 => Update { l: u32, r: u32, bc: (M, M) }
13        1 => Fold { l: u32, r: u32 }
14    }
15}
16
17#[verify::library_checker("range_affine_range_sum_large_array")]
18pub fn range_affine_range_sum_large_array(reader: impl Read, writer: impl Write) {
19    prepare_io!(reader, writer);
20    sc!(n: u32, q, queries: [Query; q]);
21    let mut values: Vec<_> = [0, n]
22        .into_iter()
23        .chain(queries.iter().flat_map(|&query| {
24            let (Query::Update { l, r, .. } | Query::Fold { l, r }) = query;
25            [l, r]
26        }))
27        .collect();
28    values.radix_sort_by_key(|&x| x);
29    values.dedup();
30    let search = StaticSearch::from_sorted(&values);
31    let mut seg = LazySegmentTree::<RangeSumRangeLinear<M>>::from_vec(
32        values
33            .windows(2)
34            .map(|w| (M::zero(), M::from(w[1] - w[0])))
35            .collect(),
36    );
37    let endpoints: Vec<_> = queries
38        .iter()
39        .flat_map(|&query| {
40            let (Query::Update { l, r, .. } | Query::Fold { l, r }) = query;
41            [l, r]
42        })
43        .collect();
44    let mut positions = vec![0; endpoints.len()];
45    search.lower_bound_batch(&endpoints, &mut positions);
46    for (query, &[l, r]) in queries.into_iter().zip(positions.as_chunks::<2>().0) {
47        match query {
48            Query::Update { bc, .. } => {
49                seg.update(l..r, bc);
50            }
51            Query::Fold { .. } => {
52                pp!(seg.fold(l..r).0);
53            }
54        }
55    }
56}