G - Sum of Binom(A, B) 解説
by
yuto1115
解説
\(A_i=j\) を満たす \(i\) の個数を \(C_a(j)\)、\(B_i=j\) を満たす \(i\) の個数を \(C_b(j)\)、数列 \(A,B\) の要素の最大値を \(K\) とおくと、
\[ \displaystyle \begin{aligned} \sum_{i=1}^{N} \sum_{j=1}^{M} \binom{A_i}{B_j} &= \sum_{i=1}^{K} \sum_{j=1}^{i} C_a(i)\cdot C_b(j)\cdot \binom{i}{j} \\ &= \sum_{i=1}^{K} \sum_{j=1}^{i} C_a(i)\cdot C_b(j)\cdot \frac{i!}{j!(i-j)!} \\ &= \sum_{i=1}^{K} \left(C_a(i)\cdot (i!)\cdot \sum_{j=1}^{i} \frac{C_b(j)}{j!(i-j)!}\right) \\ &= \sum_{i=1}^{K} \left(C_a(i)\cdot (i!)\cdot \sum_{j+k=i} \left(\frac{C_b(j)}{j!}\cdot \frac{1}{k!}\right)\right) \\ \end{aligned} \]
となります。
したがって、\(\displaystyle f(x)=\sum_{j=0}^{K} \frac{C_b(j)}{j!}x^j, \ \ g(x)=\sum_{j=0}^{K} \frac{x^j}{j!}\) と定義し、\(h(x)=f(x)\cdot g(x)\) とおくと、求める値は \(\displaystyle \sum_{i=1}^{K} C_a(i)\cdot (i!)\cdot[x^i]h\) となります。
よって、数論変換 (NTT) を用いて \(h(x)\) を計算することで、本問題を \(O(N+M+K\log K)\) で解くことができます。実装においては AC Library (Convolution) を利用することができます。
実装例 (C++) :
#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
const int C = 500010;
mint fac[C + 1];
mint ifac[C + 1];
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
fac[0] = 1;
for (int i = 0; i < C; i++) fac[i + 1] = fac[i] * (i + 1);
ifac[C] = fac[C].inv();
for (int i = C - 1; i >= 0; i--) ifac[i] = ifac[i + 1] * (i + 1);
int n, m;
cin >> n >> m;
vector<int> ca(C), cb(C);
for (int i = 0; i < n; i++) {
int a;
cin >> a;
++ca[a];
}
for (int i = 0; i < m; i++) {
int b;
cin >> b;
++cb[b];
}
vector<mint> f(C), g(C);
for (int i = 0; i < C; i++) {
f[i] = ifac[i] * cb[i];
g[i] = ifac[i];
}
auto h = convolution(f, g);
mint ans = 0;
for (int i = 0; i < C; i++) {
ans += h[i] * fac[i] * ca[i];
}
cout << ans.val() << endl;
}
投稿日時:
最終更新: