Official

D - ボールの箱詰め / Packing Balls into Boxes Editorial by sounansya


与えられる \(U_i,V_i\) をグラフとして見た時の連結成分数を \(X\) とします。すると、各連結成分は全て同じ箱に入れる必要があるため、求める答えは以下の問題の答えと同じになります:

\(N\) 個の区別できる箱に \(X\) 個の区別できるボールを入れる方法は何通りか? \(998244353\) で割ったあまりを求めよ。ただし、全ての箱に \(1\) つ以上のボールが入っている必要がある。

この問題を包除原理を用いて解くことを考えます。

\(S \subset \lbrace 1,2,\ldots,N\rbrace\) を固定し、\(i\in S\) を満たす \(i\) に対し箱 \(i\) にボールを入れないことを考えます。

残った箱の候補は \(N-|S|\) 個で、これらに \(X\) 個のボールを入れる方法は \(\displaystyle (N-|S|)^X\) 通りあります。この値は \(S\) の要素に依存せず、\(S\) の要素数によってのみ決まります。

したがって、包除原理から求める値は \(\displaystyle \sum_{i=0}^N (-1)^i \binom Ni (N-i)^X\) を \(998244353\) で割ったあまりとなります。

二項係数は階乗とその逆元を事前に計算しておくことで高速に計算しておくことができます。計算量は \(O(N\log K)\) ですが、\(x=1,2,\ldots,N\) に対し \(x^X\) をまとめて計算することで \(O(N)\) 時間に改善することもできます。

実装例

from atcoder import dsu

n, k, m = map(int, input().split())
d = dsu.DSU(k)
for i in range(m):
    u, v = map(int, input().split())
    d.merge(u - 1, v - 1)
k = len(d.groups())
MOD = 998244353
fact = [1] * (n + 1)
finv = [1] * (n + 1)
for i in range(1, n + 1):
    fact[i] = i * fact[i - 1] % MOD
finv[n] = pow(fact[n], MOD - 2, MOD)
for i in range(n, 0, -1):
    finv[i - 1] = i * finv[i] % MOD

def comb(n, k):
    if n < 0 or k < 0 or n - k < 0:
        return 0
    return fact[n] * finv[k] % MOD * finv[n - k] % MOD

ans = 0
for i in range(n + 1):
    ans += pow(-1, i, MOD) * comb(n, i) % MOD * pow(n - i, k, MOD) % MOD
print(ans % MOD)