library_checker/data_structure/
rectangle_sum.rs1use 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}