Skip to main content

library_checker/data_structure/
point_add_rectangle_sum.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation,
4    algorithm::SliceSortExt,
5    data_structure::{CompressedSegmentTree2d, WaveletMatrix, WaveletMatrixPointAdd},
6};
7
8competitive::define_enum_scan! {
9    #[derive(Clone, Copy)]
10    enum Query: u8 {
11        0 => Add { x: u32, y: u32, w: u64 }
12        1 => Sum { l: u32, d: u32, r: u32, u: u32 }
13    }
14}
15
16#[verify::library_checker("point_add_rectangle_sum")]
17pub fn point_add_rectangle_sum(reader: impl Read, writer: impl Write) {
18    prepare_io!(reader, writer);
19    sc!(n, q, xyw: [(u32, u32, u64); iter n]);
20    let mut points: Vec<_> = xyw.map(|(x, y, w)| (x, y, w as i64)).collect();
21    sc!(queries: [Query; q]);
22    points.extend(queries.iter().filter_map(|&query| match query {
23        Query::Add { x, y, .. } => Some((x, y, 0)),
24        Query::Sum { .. } => None,
25    }));
26    let mut order: Vec<_> = (0..points.len()).collect();
27    order.radix_sort_by_key(|&i| points[i].0);
28    let mut positions = vec![0; points.len()];
29    let mut xs = Vec::with_capacity(points.len());
30    let mut ys = Vec::with_capacity(points.len());
31    let mut weights = Vec::with_capacity(points.len());
32    for (i, &point) in order.iter().enumerate() {
33        positions[point] = i;
34        let (x, y, w) = points[point];
35        xs.push(x);
36        ys.push(y);
37        weights.push(w);
38    }
39    let wm = WaveletMatrix::new(ys);
40    let mut fold: WaveletMatrixPointAdd<_, AdditiveOperation<i64>> = wm.build_point_add(&weights);
41
42    let mut point = n;
43    for query in queries {
44        match query {
45            Query::Add { w, .. } => {
46                fold.update(positions[point], w as i64);
47                point += 1;
48            }
49            Query::Sum { l, d, r, u } => {
50                let l = xs.partition_point(|&x| x < l);
51                let r = xs.partition_point(|&x| x < r);
52                pp!(fold.fold_range(d..u, l..r));
53            }
54        }
55    }
56}
57
58#[verify::library_checker("point_add_rectangle_sum")]
59pub fn point_add_rectangle_sum_compressed_segment_tree(reader: impl Read, writer: impl Write) {
60    prepare_io!(reader, writer);
61    sc!(n, q, xyw: [(u32, u32, u64); n], queries: [Query; q]);
62    let points: Vec<_> = xyw
63        .iter()
64        .map(|&(x, y, _)| (x, (y,)))
65        .chain(queries.iter().filter_map(|&query| {
66            if let Query::Add { x, y, .. } = query {
67                Some((x, (y,)))
68            } else {
69                None
70            }
71        }))
72        .collect();
73
74    let mut seg = CompressedSegmentTree2d::<AdditiveOperation<u64>, _, _>::new(&points);
75    for &(x, y, w) in &xyw {
76        seg.update(&(x, (y,)), &w);
77    }
78
79    for query in queries {
80        match query {
81            Query::Add { x, y, w } => {
82                seg.update(&(x, (y,)), &w);
83            }
84            Query::Sum { l, d, r, u } => {
85                let ans = seg.fold(&(l..r, (d..u,)));
86                pp!(ans);
87            }
88        }
89    }
90}