公式

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;
}

投稿日時:
最終更新: