POJ 1990 : MooFest - 忘れても大丈夫

問題. MooFest

1次元上の異なる位置に $n$ 頭の牛がいる． $i$ 番目の牛は聴力に関する閾値 $v_i$ を持っている． $i$ 番目と $j$ 番目の牛が互いに会話するためには $\max\{v_i, v_j\} \times \left|x_i - x_j \right|$ 以上の音量で話す必要がある．ここで， $x_i$ と $x_j$ はそれぞれ $i$ 番目と $j$ 番目の座標を表す．すべての牛が互いに最低限の音量で会話したときの音量の総和を求めよ．

制約： $1 \le n \le 20,000$ , 　 $1 \le v_i \le 20,000$ ,　 $1 \le x_i \le 20,000$

解法. BIT

$n \le 20,000$ のため $O(n^2)$ の愚直な全探索では間に合わないのでより高速な解法が必要となる．
閾値で昇順にソートした牛の列を $C = (c_1, c_2, \ldots, c_n)$ とする．すなわち， $i < j$ ならば $c_i.v < c_j.v$ とする．また， $C_i = (c_1, \ldots, c_i)$ とする．ここで， $c_{i+1}$ と $c_k \in C_i$ との間で必要な音量は，
　 $max\{c_{i+1}.v, c_k.v\} \times \left| c_{i+1}.x - c_k.x \right| = c_{i+1}.v \times \left| c_{i+1}.x - c_k.x \right|$
となる．よって， $c_{i+1}$ と $C_i$ 間の音量の総和は，
　 $c_{i+1}.v \times \sum_{k=1}^{i} \left| c_{i+1}.x - c_k.x \right|$
となるので，各 $i$ で $\sum_{k=1}^1 \left| c_{i+1}.x - c_k.x \right|$ が高速に計算出来れば嬉しい．
$C_i$ を次のように2つに分割する．

$C_{i1} = \{c_k \mid c_k.x < c_{i+1}.x\}$
$C_{i2} = \{c_k \mid c_{i+1}.x < c_k.x\}$

このとき， $c_{i+1}$ と $C_i$ 間の距離の総和は次のようになる．
　 $\sum_{k=1}^{i} \left| c_{i+1}.x - c_k.x \right| = \left(c_{i+1}.x \times \left|C_{i1}\right| - \sum_{c_k \in C_{i1}} c_k.x \right)$
　　　　　　　　　　　　　　 $+ \left(\sum_{c_k \in C_{i2}} c_k.x - c_{i+1}.x \times \left|C_{i2}\right| \right)$
Binary Indexed Tree (BIT) [1] を2つ用いると上の式を $O(\log i)$ 時間で計算することができる． $\left|C_{i1}\right|$ と $\left| C_{i2} \right|$ を求めるBITの要素 $b_x$ の値は距離 $x$ の牛が登録されていれば1，されていなければ0とする．次に， $\sum_{c_k \in C_{i1}} c_k.x$ と $\sum_{c_k \in C_{i2}} c_k.x$ を求めるBITの要素 $b_x$ の値は距離 $x$ にいる牛の閾値とする．

計算時間： $O(n \log x_{\max})$ 　（ $x_{\max} := \max_{c_k \in C} c_k.x$ ）

ソースコードを表示

#include <iostream>
#include <vector>
#include <algorithm>
 
using namespace std;
 
template<class T>
struct BinaryIndexedTree {
private:
    vector<T> bit;
    int size;
 
public:
    BinaryIndexedTree(int num) : bit(num + 1, 0), size(num) {}
 
    T Sum(int a) { // bit[1] + ... + bit[a]  : O(log N)
        T ret = 0;
        for (int x = a; x > 0; x -= x & -x) ret += bit[x];
        return ret;
    }
    T Sum(const int l, const int r) { return Sum(r) - Sum(l - 1); }
    void Add(int a, int w) {    // bit[a] += w  :  O(log N)
        for (int x = a; x <= size; x += x & -x) bit[x] += w;
    }
};
 
 
int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
 
    int n, max_x = 0;
 
    cin >> n;
    vector<pair<int, int> > cow(n);
    for (int i = 0; i < n; ++i) {
        cin >> cow[i].first >> cow[i].second;
        max_x = max(max_x, cow[i].second);
    }
 
    sort(cow.begin(), cow.end());
 
    long long sum = 0;
    BinaryIndexedTree<int> bit_num(max_x);
    BinaryIndexedTree<long long> bit_x(max_x);
    for (int i = 0; i < n; ++i) {
        const long long x = cow[i].second, v = cow[i].first;
 
        sum += v * ((bit_num.Sum(0, x) * x - bit_x.Sum(0, x)) + (bit_x.Sum(x, max_x) - bit_num.Sum(x, max_x) * x));
        bit_num.Add(x, 1);
        bit_x.Add(x, x);
    }
 
    cout << sum << endl;
 
    return 0;
}