G - Sum of Binom(A, B) 解説 by en_translator
Let \(C_a(j)\) be the number of indices \(i\) with \(A_i=j\), \(C_b(j)\) be the number of indices \(i\) with \(B_i=j\), and \(K\) be the maximum value among the elements of \(A\) and \(B\). Then we have
\[ \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} \]
Therefore, defining \(\displaystyle f(x)=\sum_{j=0}^{K} \frac{C_b(j)}{j!}x^j, \ \ g(x)=\sum_{j=0}^{K} \frac{x^j}{j!}\) and \(h(x)=f(x)\cdot g(x)\), the sought value can be represented as \(\displaystyle \sum_{i=1}^{K} C_a(i)\cdot (i!)\cdot[x^i]h\).
Hence, by computing \(h(x)\) using Numeric Theory Transform (NTT), the problem can be solved in \(O(N+M+K\log K)\) time. For implementation, one can use AC Library (Convolution).
Sample code (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;
}
投稿日時:
最終更新: