Codeforces Round #404 (Div.2) E. Anton and Permutation

問題文

http://codeforces.com/contest/785/problem/E

解法

平方分割 + 2次元BITで解いた．[0...m] x [0...n] で二次元．m は n / sqrt(n)．
l[i] から r[i] の間に完全に含まれているブロックは BIT 上で計算し，端っこの部分は生で持っている配列を使って計算する．
bit.sum(lb, a, rb, b) とすると lb から rb - 1 までのブロックにおける [a..b) の数を求められる．

計算量は O(Q * (sqrt(N) + log^2(N)))．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

template <typename T>
class fenwick_tree2d {
public:
    fenwick_tree2d(int h, int w)
        : H(h), W(w), dat(H, std::vector<T>(W))
    {}

    void add(int i, int j, T value) {
        for(; i < H; i |= i + 1) {
            int jj = j;
            for(; jj < W; jj |= jj + 1) {
                dat[i][jj] += value;
            }
        }
    }

    // [0..i] x [0..j]
    T sum(int i, int j) const {
        T res = 0;
        for(; i >= 0; i = (i & (i + 1)) - 1) {
            int jj = j;
            for(; jj >= 0; jj = (jj & (jj + 1)) - 1) {
                res += dat[i][jj];
            }
        }
        return res;
    }

    // [i1...i2) x [j1...j2)
    T sum(int i1, int j1, int i2, int j2) const {
        return sum(i2 - 1, j2 - 1) - sum(i1 - 1, j2 - 1) - sum(i2 - 1, j1 - 1) + sum(i1 - 1, j1 - 1);
    }

private:
    int const H, W;
    std::vector<std::vector<T>> dat;
};

int main() {
    int n, q;
    cin >> n >> q;
    vector<int> l(q), r(q);
    for(int i = 0; i < q; ++i) {
        cin >> l[i] >> r[i];
        l[i]--; r[i]--;
        if(l[i] > r[i]) {
            swap(l[i], r[i]);
        }
    }

    vector<int> a(n);
    for(int i = 0; i < n; ++i) {
        a[i] = i;
    }
    int const B = sqrt(n) + 1;
    fenwick_tree2d<int> bit(n / B + 1, n);
    for(int i = 0; i < n; ++i) {
        bit.add(i / B, i, 1);
    }

    ll inv = 0;
    for(int i = 0; i < q; ++i) {
        if(l[i] != r[i]) {
            int lb = l[i] / B, rb = r[i] / B;
            if(lb + 1 < rb) {
                int l_lt = bit.sum(lb + 1, 0, rb, a[l[i]]);
                int l_gt = bit.sum(lb + 1, a[l[i]] + 1, rb, n);
                int r_lt = bit.sum(lb + 1, 0, rb, a[r[i]]);
                int r_gt = bit.sum(lb + 1, a[r[i]] + 1, rb, n);
                inv += l_gt - l_lt + r_lt - r_gt;
            }
            if(lb == rb) {
                for(int j = l[i] + 1; j < r[i]; ++j) {
                    inv += (a[l[i]] < a[j] ? 1 : -1) + (a[r[i]] < a[j] ? -1 : 1);
                }
            } else {
                for(int j = l[i] + 1; j < (lb + 1) * B; ++j) {
                    inv += (a[l[i]] < a[j] ? 1 : -1) + (a[r[i]] < a[j] ? -1 : 1);
                }
                for(int j = rb * B; j < r[i]; ++j) {
                    inv += (a[l[i]] < a[j] ? 1 : -1) + (a[r[i]] < a[j] ? -1 : 1);
                }
            }
            inv += (a[l[i]] < a[r[i]] ? 1 : -1);

            bit.add(l[i] / B, a[l[i]], -1);
            bit.add(r[i] / B, a[r[i]], -1);
            swap(a[l[i]], a[r[i]]);
            bit.add(l[i] / B, a[l[i]], 1);
            bit.add(r[i] / B, a[r[i]], 1);
        }
        cout << inv << endl;
    }
}

感想

メモリもTLEも怖すぎる．