-
Notifications
You must be signed in to change notification settings - Fork 16
/
rectangle-add-rectangle-sum.hpp
90 lines (79 loc) · 2.66 KB
/
rectangle-add-rectangle-sum.hpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
#pragma once
#include <vector>
using namespace std;
#include "../data-structure/binary-indexed-tree.hpp"
// https://hitonanode.github.io/cplib-cpp/data_structure/rectangle_add_rectangle_sum.hpp
template <class I, class T>
class RectangleAddRectangleSum {
struct AddQuery {
I xl, xr, yl, yr;
T val;
};
struct SumQuery {
I xl, xr, yl, yr;
};
vector<AddQuery> add_queries;
vector<SumQuery> sum_queries;
public:
RectangleAddRectangleSum() = default;
void add_rectangle(I xl, I xr, I yl, I yr, T val) {
add_queries.push_back(AddQuery{xl, xr, yl, yr, val});
}
void add_query(I xl, I xr, I yl, I yr) {
sum_queries.push_back(SumQuery{xl, xr, yl, yr});
}
vector<T> solve() const {
vector<I> ys;
for (const auto &a : add_queries) {
ys.push_back(a.yl);
ys.push_back(a.yr);
}
sort(ys.begin(), ys.end());
ys.erase(unique(ys.begin(), ys.end()), ys.end());
const int Y = ys.size();
vector<tuple<I, int, int>> ops;
for (int q = 0; q < int(sum_queries.size()); ++q) {
ops.emplace_back(sum_queries[q].xl, 0, q);
ops.emplace_back(sum_queries[q].xr, 1, q);
}
for (int q = 0; q < int(add_queries.size()); ++q) {
ops.emplace_back(add_queries[q].xl, 2, q);
ops.emplace_back(add_queries[q].xr, 3, q);
}
sort(ops.begin(), ops.end());
BinaryIndexedTree<T> b00(Y), b01(Y), b10(Y), b11(Y);
vector<T> ret(sum_queries.size());
for (auto o : ops) {
int qtype = get<1>(o), q = get<2>(o);
if (qtype >= 2) {
const AddQuery &query = add_queries.at(q);
int i = lower_bound(ys.begin(), ys.end(), query.yl) - ys.begin();
int j = lower_bound(ys.begin(), ys.end(), query.yr) - ys.begin();
T x = get<0>(o);
T yi = query.yl, yj = query.yr;
if (qtype & 1) swap(i, j), swap(yi, yj);
b00.add(i, x * yi * query.val);
b01.add(i, -x * query.val);
b10.add(i, -yi * query.val);
b11.add(i, query.val);
b00.add(j, -x * yj * query.val);
b01.add(j, x * query.val);
b10.add(j, yj * query.val);
b11.add(j, -query.val);
} else {
const SumQuery &query = sum_queries.at(q);
int i = lower_bound(ys.begin(), ys.end(), query.yl) - ys.begin();
int j = lower_bound(ys.begin(), ys.end(), query.yr) - ys.begin();
T x = get<0>(o);
T yi = query.yl, yj = query.yr;
if (qtype & 1) swap(i, j), swap(yi, yj);
i--, j--;
ret[q] +=
b00.sum(i) + b01.sum(i) * yi + b10.sum(i) * x + b11.sum(i) * x * yi;
ret[q] -=
b00.sum(j) + b01.sum(j) * yj + b10.sum(j) * x + b11.sum(j) * x * yj;
}
}
return ret;
}
};