IOI '07 P5 - Pairs

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int MX1 = 75000000, MX2 = 75000, MX3 = 75;

// Fenwick Tree
void update(vector<int> &fenwick, int index, int value) {
    while (index < (int)fenwick.size()) {
        fenwick[index] += value;
        index += index & -index;
    }
}
int getSum(const vector<int> &fenwick, int index) {
    int sum = 0;
    while (index) {
        sum += fenwick[index];
        index -= index & -index;
    }
    return sum;
}
int rangeSum(const vector<int> &fenwick, int l, int r) {
    return getSum(fenwick, min(r, (int)fenwick.size()-1)) - getSum(fenwick, max(l-1, 0));
}

int main() {
    int b, n, d, m;
    cin >> b >> n >> d >> m;
    long long pairs = 0;
    if (b == 1) {
        vector<int> x(n);
        for (int &a : x) cin >> a;
        sort(x.begin(), x.end());
        for (int i = 0, j = 0; i < n; ++i) {
            while (x[j]+d < x[i]) ++j;
            pairs += i-j;
        }
    } else if (b == 2) {
        // convert from manhattan distance to king's distance
        vector<pair<int, int>> a(n); // x-y, x+y
        for (int i = 0; i < n; ++i) {
            int x, y; cin >> x >> y;
            a[i] = {x-y+MX2, x+y}; // add MX2 to keep positive
        }
        // each new coordinate is between 1 and 2*MX2
        sort(a.begin(), a.end()); // sort by x

        // Sliding window
        vector<int> fenwick(1+2*MX2); // frequency by y position
        for (int i = 0, j = 0; i < n; update(fenwick, a[i++].second, +1)) {
            while (a[j].first+d < a[i].first) {
                update(fenwick, a[j++].second, -1);
            }
            pairs += rangeSum(fenwick, a[i].second-d, a[i].second+d);
        }
    } else { // b == 3
        // convert (x, y) to (x-y, x+y) - from manhattan to king's distance
        vector<int> x(n), y(n), z(n);
        for (int i = 0; i < n; ++i) {
            cin >> x[i] >> y[i] >> z[i];
            y[i] += x[i]; // x+y
            x[i] = 2*x[i] - y[i] + MX3; // x-y (add MX3 to keep positive)
        }

        // prefix area sums (can also be done with sliding windows between pairs of layers)
        vector<vector<vector<int>>> f(1+MX3, vector<vector<int>>(1+2*MX3, vector<int>(1+2*MX3)));
        for (int i = 0; i < n; ++i) ++f[z[i]][x[i]][y[i]];
        for (int k = 1; k <= MX3; ++k) {
            for (int i = 1; i <= 2*MX3; ++i) {
                for (int j = 1; j <= 2*MX3; ++j) {
                    f[k][i][j] += f[k][i-1][j] + f[k][i][j-1] - f[k][i-1][j-1];
                }
            }
        }
        auto rectangleSum = [&f](int z, int x1, int y1, int x2, int y2) {
            if (x1 < 1) x1 = 1;
            if (y1 < 1) y1 = 1;
            if (x2 > 2*MX3) x2 = 2*MX3;
            if (y2 > 2*MX3) y2 = 2*MX3;
            return f[z][x2][y2] - f[z][x1-1][y2] - f[z][x2][y1-1] + f[z][x1-1][y1-1];
        };
        for (int i = 0; i < n; ++i) {
            for (int k = 1; k <= MX3; ++k) {
                int d2 = d - abs(k - z[i]);
                if (d2 < 0) continue;
                pairs += rectangleSum(k, x[i]-d2, y[i]-d2, x[i]+d2, y[i]+d2);
            }
        }
        pairs = (pairs-n)/2; // remove self-pairs and half since double-counted
    }
    cout << pairs << endl;
    return 0;
}