Official

F - Merge Sets Editorial by yuto1115

解説

まず、\(f(A,B)\) の定義は以下のように言い換えられます。

  • \(\displaystyle f(A,B) = \sum_{i \in A} (A \cup B \text{ の中に含まれる } i\text{ 以下の数の個数})\)

これを用いると、求める答えは、

\[ \begin{aligned} \sum_{1 \leq i < j \leq N} f(S_i,S_j) &= \sum_{1 \leq i < j \leq N} \sum_{k=1}^{M} (S_i \cup S_j \text{ の中に含まれる } A_{i,k}\text{ 以下の数の個数})\\ &= \sum_{1 \leq i < j \leq N} \left( \sum_{k=1}^{M} (S_i \text{ の中に含まれる } A_{i,k}\text{ 以下の数の個数}) + \sum_{k=1}^{M} (S_j \text{ の中に含まれる } A_{i,k}\text{ 以下の数の個数}) \right)\\ &= \sum_{1 \leq i < j \leq N} \left( \frac{M(M+1)}{2} + \sum_{k=1}^{M} (S_j \text{ の中に含まれる } A_{i,k}\text{ 以下の数の個数}) \right)\\ &= \frac{M(M+1)}{2} \cdot \frac{N(N-1)}{2} + \sum_{1 \leq i < j \leq N}\sum_{k=1}^{M} (S_j \text{ の中に含まれる } A_{i,k}\text{ 以下の数の個数})\\ &= \frac{M(M+1)}{2} \cdot \frac{N(N-1)}{2} + \sum_{i=1}^{N}\sum_{k=1}^{M} (S_{i+1}\cup \dots \cup S_N \text{ の中に含まれる } A_{i,k}\text{ 以下の数の個数})\\ \end{aligned} \]

というように変形できます。後半の \(2\) 重シグマの部分については、\(i\) の降順に \(A_{i,j}\) を走査しつつ fenwick tree 等を用いることによって \(O(NM \log NM)\) で求めることができます。なお、最初に \(A\) を座標圧縮する必要があることに注意してください。

実装例 (C++) :

#include<bits/stdc++.h>
#include<atcoder/fenwicktree>

using namespace std;
using namespace atcoder;

using ll = long long;

int main() {
    int n, m;
    cin >> n >> m;
    vector a(n, vector<int>(m));
    vector<int> v;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            cin >> a[i][j];
            v.push_back(a[i][j]);
        }
    }
    sort(v.begin(), v.end());
    
    ll ans = (ll) m * (m + 1) / 2 * n * (n - 1) / 2;
    fenwick_tree<int> fw(n * m);
    for (int i = n - 1; i >= 0; i--) {
        for (int j = 0; j < m; j++) {
            a[i][j] = lower_bound(v.begin(), v.end(), a[i][j]) - v.begin();
            ans += fw.sum(0, a[i][j]);
        }
        for (int j = 0; j < m; j++) {
            fw.add(a[i][j], 1);
        }
    }
    cout << ans << endl;
}

posted:
last update: