haar_lib/ds/
static_rectangle_add_sum.rs

1//! 矩形加算矩形総和
2//!
3//! # References
4//! - <https://ei1333.hateblo.jp/entry/2022/06/10/022355>
5//!
6//! # Problems
7//! - <https://judge.yosupo.jp/problem/static_rectangle_add_rectangle_sum>
8
9use std::ops::{Add, AddAssign, Mul, Neg, Sub};
10
11use crate::{algo::compressor::CompressorBuilder, ds::fenwick_add::*, num::one_zero::*};
12
13struct QAdd<T>(i64, usize, T);
14struct QSum<T>(i64, usize, T, usize);
15
16/// 矩形加算矩形総和
17///
18/// すべての矩形加算クエリを行ってから、矩形総和クエリに答える。
19#[derive(Clone, Default)]
20pub struct StaticRectangleAddSum<T> {
21    add: Vec<(i64, i64, i64, i64, T)>,
22    sum: Vec<(i64, i64, i64, i64)>,
23}
24
25impl<T> StaticRectangleAddSum<T>
26where
27    T: Copy
28        + Mul<Output = T>
29        + Add<Output = T>
30        + Sub<Output = T>
31        + Neg<Output = T>
32        + AddAssign
33        + Zero
34        + One
35        + From<i64>,
36{
37    /// 空の`StaticRectangleAddSum`を返す。
38    pub fn new() -> Self {
39        Self {
40            add: vec![],
41            sum: vec![],
42        }
43    }
44
45    /// 矩形加算クエリを追加する。
46    pub fn query_add(&mut self, l: i64, d: i64, r: i64, u: i64, w: T) {
47        self.add.push((l, d, r, u, w));
48    }
49
50    /// 矩形総和クエリを追加する。
51    pub fn query_sum(&mut self, l: i64, d: i64, r: i64, u: i64) {
52        self.sum.push((l, d, r, u));
53    }
54
55    /// 矩形総和クエリの結果を返す。
56    pub fn solve(self) -> Vec<T> {
57        let mut cp = CompressorBuilder::new();
58        cp.extend(self.add.iter().flat_map(|(_, d, _, u, _)| vec![*d, *u]));
59        cp.extend(self.sum.iter().flat_map(|(_, d, _, u)| vec![*d, *u]));
60        let cp = cp.build();
61
62        let mut add = self
63            .add
64            .into_iter()
65            .flat_map(|(l, d, r, u, w)| {
66                let d = cp.index(&d);
67                let u = cp.index(&u);
68                vec![QAdd(l, d, w), QAdd(r, u, w), QAdd(l, u, -w), QAdd(r, d, -w)]
69            })
70            .collect::<Vec<_>>();
71
72        add.sort_by_key(|q| q.0);
73        add.reverse();
74
75        let mut sum = self
76            .sum
77            .iter()
78            .enumerate()
79            .flat_map(|(i, &(l, d, r, u))| {
80                let d = cp.index(&d);
81                let u = cp.index(&u);
82                vec![
83                    QSum(l, d, T::one(), i),
84                    QSum(r, u, T::one(), i),
85                    QSum(l, u, -T::one(), i),
86                    QSum(r, d, -T::one(), i),
87                ]
88            })
89            .collect::<Vec<_>>();
90
91        sum.sort_by_key(|q| q.0);
92        sum.reverse();
93
94        let mut fxy = FenwickTreeAdd::new(cp.size());
95        let mut fx = FenwickTreeAdd::new(cp.size());
96        let mut fy = FenwickTreeAdd::new(cp.size());
97        let mut f = FenwickTreeAdd::new(cp.size());
98        let mut ret = vec![T::zero(); self.sum.len()];
99
100        loop {
101            let is_sum = match (add.last(), sum.last()) {
102                (None, None) => break,
103                (Some(_), None) => false,
104                (None, Some(_)) => true,
105                (Some(a), Some(s)) => a.0 >= s.0,
106            };
107
108            if is_sum {
109                let QSum(r, y, c, index) = sum.pop().unwrap();
110                let u = *cp.get(y);
111
112                let ans = fxy.fold_to(..y)
113                    + f.fold_to(..y) * r.into() * u.into()
114                    + fy.fold_to(..y) * r.into()
115                    + fx.fold_to(..y) * u.into();
116
117                ret[index] += c * ans;
118            } else {
119                let QAdd(l, y, w) = add.pop().unwrap();
120                let d = *cp.get(y);
121
122                fxy.add(y, w * l.into() * d.into());
123                f.add(y, w);
124                fy.sub(y, w * d.into());
125                fx.sub(y, w * l.into());
126            }
127        }
128
129        ret
130    }
131}
haar_lib/ds/static_rectangle_add_sum.rs

haar_lib/ds/
static_rectangle_add_sum.rs
