# IOI '07 P5 - Pairs

Jun 13th, 2022
596
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <iostream>
2. #include <vector>
3. #include <algorithm>
4.
5. using namespace std;
6.
7. const int MX1 = 75000000, MX2 = 75000, MX3 = 75;
8.
9. // Fenwick Tree
10. void update(vector<int> &fenwick, int index, int value) {
11.     while (index < (int)fenwick.size()) {
12.         fenwick[index] += value;
13.         index += index & -index;
14.     }
15. }
16. int getSum(const vector<int> &fenwick, int index) {
17.     int sum = 0;
18.     while (index) {
19.         sum += fenwick[index];
20.         index -= index & -index;
21.     }
22.     return sum;
23. }
24. int rangeSum(const vector<int> &fenwick, int l, int r) {
25.     return getSum(fenwick, min(r, (int)fenwick.size()-1)) - getSum(fenwick, max(l-1, 0));
26. }
27.
28. int main() {
29.     int b, n, d, m;
30.     cin >> b >> n >> d >> m;
31.     long long pairs = 0;
32.     if (b == 1) {
33.         vector<int> x(n);
34.         for (int &a : x) cin >> a;
35.         sort(x.begin(), x.end());
36.         for (int i = 0, j = 0; i < n; ++i) {
37.             while (x[j]+d < x[i]) ++j;
38.             pairs += i-j;
39.         }
40.     } else if (b == 2) {
41.         // convert from manhattan distance to king's distance
42.         vector<pair<int, int>> a(n); // x-y, x+y
43.         for (int i = 0; i < n; ++i) {
44.             int x, y; cin >> x >> y;
45.             a[i] = {x-y+MX2, x+y}; // add MX2 to keep positive
46.         }
47.         // each new coordinate is between 1 and 2*MX2
48.         sort(a.begin(), a.end()); // sort by x
49.
50.         // Sliding window
51.         vector<int> fenwick(1+2*MX2); // frequency by y position
52.         for (int i = 0, j = 0; i < n; update(fenwick, a[i++].second, +1)) {
53.             while (a[j].first+d < a[i].first) {
54.                 update(fenwick, a[j++].second, -1);
55.             }
56.             pairs += rangeSum(fenwick, a[i].second-d, a[i].second+d);
57.         }
58.     } else { // b == 3
59.         // convert (x, y) to (x-y, x+y) - from manhattan to king's distance
60.         vector<int> x(n), y(n), z(n);
61.         for (int i = 0; i < n; ++i) {
62.             cin >> x[i] >> y[i] >> z[i];
63.             y[i] += x[i]; // x+y
64.             x[i] = 2*x[i] - y[i] + MX3; // x-y (add MX3 to keep positive)
65.         }
66.
67.         // prefix area sums (can also be done with sliding windows between pairs of layers)
68.         vector<vector<vector<int>>> f(1+MX3, vector<vector<int>>(1+2*MX3, vector<int>(1+2*MX3)));
69.         for (int i = 0; i < n; ++i) ++f[z[i]][x[i]][y[i]];
70.         for (int k = 1; k <= MX3; ++k) {
71.             for (int i = 1; i <= 2*MX3; ++i) {
72.                 for (int j = 1; j <= 2*MX3; ++j) {
73.                     f[k][i][j] += f[k][i-1][j] + f[k][i][j-1] - f[k][i-1][j-1];
74.                 }
75.             }
76.         }
77.         auto rectangleSum = [&f](int z, int x1, int y1, int x2, int y2) {
78.             if (x1 < 1) x1 = 1;
79.             if (y1 < 1) y1 = 1;
80.             if (x2 > 2*MX3) x2 = 2*MX3;
81.             if (y2 > 2*MX3) y2 = 2*MX3;
82.             return f[z][x2][y2] - f[z][x1-1][y2] - f[z][x2][y1-1] + f[z][x1-1][y1-1];
83.         };
84.         for (int i = 0; i < n; ++i) {
85.             for (int k = 1; k <= MX3; ++k) {
86.                 int d2 = d - abs(k - z[i]);
87.                 if (d2 < 0) continue;
88.                 pairs += rectangleSum(k, x[i]-d2, y[i]-d2, x[i]+d2, y[i]+d2);
89.             }
90.         }
91.         pairs = (pairs-n)/2; // remove self-pairs and half since double-counted
92.     }
93.     cout << pairs << endl;
94.     return 0;
95. }