POJ 1990 - MooFest - そーすにっき

久々のソースコード投稿になりますん

問題：

一次元直線上に一列に並んだ牛に座標x_i,値v_iが与えられる

$\sum_{i \in N, j \in N} max(v_i,v_j) | x_i - x_j |$

の値を求めよ.

問題の難しい点：

簡単に見えて制約が厳しい.

$|N| = 20000,x \ in {1,2,\cdots,20000}(\Leftrightarrow |x| = 20000 = |N|), time = 1000ms$

この制約上＋POJでは $O(|N|^2)$ が許されない.

そこでどうにかしてやる必要がある.

解法：

色々方法はあると思うけどとりあえずセグ木でRange Sum Queryをすることにした.

$\sum_{i \in N, j \in N} max(v_i,v_j) | x_i - x_j |$

は $v_i < v_j$ の下で

$\sum_{i \in N, j \in N,v_i < v_j} v_j | x_i - x_j |$

であるからとりあえず $v$ で昇順ソートして以下 $i < j$ とする.

次に $x_i < x_j$ であるiの集合を $L_j$ , 補集合を $R_j$ とおく.すると

$\sum_{j \in N} \left( v_j \left( |L_j| x_j - \sum_{i \in N, j \in N, i \in L_j} x_i \right) + v_j \left( \sum_{i \in N, j \in N, i \in R_j} x_i - (|R_j|) x_j \right) \right)$

$|L_j|,\sum_{i \in N, j \in N, i \in L_j} x_i$ はそれぞれRSQにデータを保持すると更新、算出共に $O(log |x|)$ で計算できる(R_j側についても同様).

これを $j \in N$ についてfor文を回すということで全体 $O(|N|log|x|)$ で計算可能.制約より $|N| = |x|$ なのでこれなら定数が重くない限り制限時間内に間に合う.

若干添字周りに $\pm 1$ が入るかもしれないが気にしない

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

using namespace std;

template <typename T>
class segtree_sum
{
  public:
    vector<T> data;
    int size;
    segtree_sum<T>(int size_) : size(1)
    {
        while (size < size_)
            size <<= 1;
        data.assign(size * 2, 0);
    }
    segtree_sum<T>(vector<T> &src) : size(1)
    {
        while (size < src.size())
            size <<= 1;
        data.assign(size * 2, 0);
        for (int i = 0; i < src.size(); i++)
        {
            update(i, src[i]);
        }
    };
    void update(int idx, int v)
    {
        int f = idx + size - 1;
        data[f] = v;
        while (f)
        {
            --f /= 2;
            data[f] = data[2 * f + 1] + data[2 * f + 2];
        }
    }
    T rsum(int l, int r)
    { //return sum[l,r)
        int f = l + size - 1;
        int k = 1;
        while (f % 2 && l + k * 2 <= r)
        {
            f /= 2;
            k *= 2;
        }
        if (l + k > r)
            return 0;
        else if (l + k == r)
            return data[f];
        else
            return data[f] + rsum(l + k, r);
    }
};

/////////////////////////////////////////////////////////////////////////

bool compare(pair<int, int> a, pair<int, int> b)
{
    return a.first < b.first;
}

#define i64 long long int

int main()
{
    int n;
    cin >> n;
    vector<pair<int, int> > moo(n);
    for (int i = 0; i < n; i++)
    {
        cin >> moo[i].first >> moo[i].second;
    }
    sort(moo.begin(), moo.end(), compare);
    segtree_sum<int> l_cows(20010);
    segtree_sum<int> l_x(20010);
    i64 res = 0LL;
    for (int i = 0; i < n; i++)
    {
        int moot = moo[i].first;
        int moox = moo[i].second;
        i64 lcows = (i64)l_cows.rsum(1, moox);
        i64 lx = (i64)l_x.rsum(1, moox);
        i64 rx = (i64)l_x.rsum(moox + 1, 20001);
        res += (i64)(lcows * moox - lx) * moot;
        res += (i64)(rx - (i - lcows) * moox) * moot;
        l_cows.update(moox, 1);
        l_x.update(moox, moox);
    }
    cout << res << endl;
}

久々に記事書いた気がする