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: