Skip to main content

library_checker/data_structure/
rectangle_sum.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation,
4    algorithm::SliceSortExt,
5    data_structure::{StaticSearch, WaveletMatrix},
6};
7
8#[verify::library_checker("rectangle_sum")]
9pub fn rectangle_sum(reader: impl Read, writer: impl Write) {
10    prepare_io!(buffered; reader, writer);
11    sc!(n, q, mut xyw: [(u32, u32, i64); n], queries: [(u32, u32, u32, u32); q]);
12    xyw.radix_sort_by_key(|&(x, ..)| x);
13    let xs: Vec<_> = xyw.iter().map(|&(x, ..)| x).collect();
14    let search = StaticSearch::from_sorted(&xs);
15    let endpoints: Vec<_> = queries.iter().flat_map(|&(l, _, r, _)| [l, r]).collect();
16    let mut positions = vec![0; endpoints.len()];
17    search.lower_bound_batch(&endpoints, &mut positions);
18    let ys = xyw.iter().map(|&(_, y, _)| y).collect();
19    let weights: Vec<_> = xyw.iter().map(|&(_, _, w)| w).collect();
20    let wm = WaveletMatrix::new(ys);
21    let fold = wm.build_fold::<AdditiveOperation<i64>>(&weights);
22    let result = fold.fold_lessthan_batch(
23        queries
24            .into_iter()
25            .zip(positions.as_chunks::<2>().0)
26            .flat_map(|((_, d, _, u), &[l, r])| [(d, l..r), (u, l..r)]),
27    );
28    for &[lower, upper] in result.as_chunks::<2>().0 {
29        pp!(upper - lower);
30    }
31}