連載しょぼんコーダー 4 先頭固定テクニック

第4回では, 先頭固定テクニックとぼくが言う, 主に定数個入力数え上げで平等性が高い状況のときに使えるDPについて書きます.

問題1

例題4-1 $N$ 個の区別する玉を $K$ 個の区別しない箱に入れる方法の数はいくつですか. ただし, 各箱には必ず $1$ 個以上の玉を入れるものとします.
つまり, 第2種Stirling数 $S(N, K)$ を求めてください.
ただし, 答えは $998244353$ で割った余りで出力してください.
時間制限：2秒、メモリ制限：1024 MB
制約

$1 \leq K \leq N \leq 300$

入力
N K

考察

$dp[i][j]$ を $i$ 個の区別する玉を $j$ 個の区別しない箱に各箱1個以上入れる方法とします.
次の $O(NK)$ のDPが知られています.

mod = 998244353
n,k = map(int,input().split())
dp = [[0] * (k+1) for j in range(n+1)]
dp[0][0] = 1

for i in range(n):
	for j in range(k):
		dp[i+1][j+1] = (j+1) * dp[i][j+1] + dp[i][j]
		dp[i+1][j+1] %= mod

print(dp[n][k])

今回は, 先頭固定テクニックで解いてみます. $dp[i][j]$ について考えます. ある1つの箱について考えます. その箱は $m$ 個の玉を含んでいるものとします. このとき, ある玉の組み合わせを固定すると, その組み合わせを含むものは, $dp[i-m][j-1]$ 通りとなります.
それはそうなのですが, 実際にすべての組み合わせについてこれを足し合わせると（つまり $\binom{i}{m}$ を掛けると）, 次の画像のように同じ結果になるものがダブルカウントしてしまうことがあります.

ダブルカウントしない方法として, 「今見ている配列のうち, 先頭を固定する」ということが考えられます. 先頭を固定すれば, 後続は $\binom{i-1}{m-1}$ 通りですが, 先頭が固定されていることにより遷移先で先頭が証人になっているので, 後々になってダブルカウントすることはありません. これが先頭固定テクニックです.
$dp[i][j]$ は, $1\leq m\leq i$ のすべてについて $dp[i-m][j-1] \times \binom{i-1}{m-1}$ の総和となります. $x!, 1/y!$ を $O(N + \log \text{MOD})$ で前計算することにより, このDPは $O(N^2K)$ となります.

ACコード

mod = 998244353
N = 400
fact = [1]*(N+1)
factinv = [1]*(N+1)

for i in range(2, N+1):
	fact[i] = fact[i-1] * i % mod

factinv[-1] = pow(fact[-1], mod-2, mod)
for i in range(N-1, 1, -1):
	factinv[i] = factinv[i+1] * (i+1) % mod

def cmb(a, b):
	if (a < b) or (b < 0):
		return 0
	return fact[a] * factinv[b] % mod * factinv[a-b] % mod

n,k = map(int,input().split())
dp = [[0] * (k+1) for j in range(n+1)]
dp[0][0] = 1

for i in range(1,n+1):
	for j in range(1,k+1):
		for m in range(1,i+1):
			dp[i][j] += dp[i-m][j-1] * cmb(i-1, m-1) % mod
			dp[i][j] %= mod

print(dp[n][k])

時間計算量はあまりよくないけれど, 有名な遷移とはまた異なる遷移ができました.
余談ですが, $\sum^{i}_{m=1} \binom{i-1}{m-1} dp[i-m][j-1] = \sum^{i-1}_{m=0} \binom{i-1}{m} dp[m][j-1]$ と変形できます. $dp[i][j] = \sum^{i-1}_{m=0} \binom{i-1}{m} dp[m][j-1] = \sum^{i-2}_{m=0} \binom{i-1}{m} dp[m][j-1] + dp[i-1][j-1]$ に注目すると, 次の式が得られます.
$\sum^{i}_{m=0} \left\{\binom{i+1}{m} - (j+1) \binom{i}{m}\right\} S(m, j) = 0$ （ただし $i, j$ は任意の非負整数）

問題2

例題4-2 ACM-ICPC Japan Alumni Group Contest 2019 2日目 H
問題文は長いので上のリンクを見てください.
時間制限：2秒、メモリ制限：256 MB
制約

$1 \leq N \leq 10^5$

考察

長さ $N$ の順列 $p$ のスコア $f(p)$ の総和を $A[N] = \sum f(p)$ , スコアの2乗の総和を $B[N] = \sum f(p)^2$ とします.
このとき, スコアの平均値は $a=\frac{1}{N!} A[N]$ , 分散は $\frac{1}{N!} \sum (f(p)^2 - 2 f(p) a + a^2) = \frac{1}{N!} (B[N] - 2 A[N] a + a^2 N!)$ です.
よって, $A[N], B[N]$ を求めればokです.

順列 $p$ を対称群 $i \mapsto p(i)$ の元だと見ると, 「サイクル森の1つの連結成分の数」は, 「巡回置換の長さ」（つまり $G(p)$ は閉路からのみなる）と確認することができます.
その閉路に注目します. $A[i]$ を求めてみましょう. $m$ を, 「先頭に含まれるグループの連結成分の数」とします. このとき, そのようなグループの決め方は $\binom{i-1}{m-1}$ 個, さらにグループの中の閉路の作り方が $(m-1)!$ 個あります. （ $m$ 次対称群の長さ $m$ の巡回置換は $(m-1)!$ 個あるので）
$A[i-m]$ から乗算されるので, 閉路の大きさ $m$ も掛けて, $m$ $A[i-m] \times \binom{i-1}{m-1} \times (m-1)! \times m$ が $A[i]$ に加算されます.
$m=1,2,\cdots, i$ を考慮して, $A[i] = \sum^{i}_{m=1} A[i-m] \times \binom{i-1}{m-1} \times (m-1)! \times m$ となります.

このように, 先頭固定テクニックは対称群（置換）と相性が良いです.

同様に $B[i] = \sum^{i}_{m=1} B[i-m] \times \binom{i-1}{m-1} \times (m-1)! \times m^2$ です.
実はこれで $O(N^2)$ コードができました. しかし間に合わないので工夫をします.

式変形をすると,
$A[i] = \sum^{i}_{m=1} A[i-m] \times \binom{i-1}{m-1} \times (m-1)! \times m$
$= \sum^{i}_{m=1} A[i-m] \times (i-1)! / (m-1)! / (i-m)! \times (m-1)! \times m$
$= (i-1)! \sum^{i}_{m=1} A[i-m] / (i-m)! \times m$
で, 順序を変えると
$A[i] = (i-1)! \sum^{i-1}_{m=0} A[m] / m! \times (i-m)$
$= (i-1)! \left\{i \times \sum^{i-1}_{m=0} A[m] / m! - \sum^{i-1}_{m=0} A[m] / m! \times m\right\}$

同様に
$A[i] = (i-1)! \sum^{i-1}_{m=0} B[m] / m! \times (i-m)^2$
$= (i-1)! \left\{i^2 \times \sum^{i-1}_{m=0} B[m] / m! - 2i \sum^{i-1}_{m=0} B[m] / m! \times m + \sum^{i-1}_{m=0} B[m] / m! \times m^2 \right\}$
です.

$\sum^{i-1}_{m=0} A[m] / m!$
$\sum^{i-1}_{m=0} A[m] / m! \times m$
$\sum^{i-1}_{m=0} B[m] / m!$
$\sum^{i-1}_{m=0} B[m] / m! \times m$
$\sum^{i-1}_{m=0} B[m] / m! \times m^2$
を持つ変数を管理すれば, この5つは, $i$ を1つ増やすごとに $O(1)$ で更新できます. よって累積和的に各 $A[i], B[i]$ を $O(1)$ で更新でき, 全体で $O(N + \log \text{MOD})$ で解けます.

ACコード

mod = 10**9 + 7
N = 100256
fact = [1]*(N+1)
factinv = [1]*(N+1)

for i in range(2, N+1):
	fact[i] = fact[i-1] * i % mod

factinv[-1] = pow(fact[-1], mod-2, mod)
for i in range(N-1, 1, -1):
	factinv[i] = factinv[i+1] * (i+1) % mod

def cmb(a, b):
	if (a < b) or (b < 0):
		return 0
	return fact[a] * factinv[b] % mod * factinv[a-b] % mod

n = int(input())
a0 = 1
a1 = 0
b0 = 1
b1 = 0
b2 = 0
t = 0
s = 0
for i in range(1, n+1):
	t = fact[i-1] * (i*a0 - a1) % mod
	s = fact[i-1] * (i*i*b0 - 2*i*b1 + b2) % mod
	a0 = (a0 + t*factinv[i]) % mod
	a1 = (a1 + t*i*factinv[i]) % mod
	b0 = (b0 + s*factinv[i]) % mod
	b1 = (b1 + s*i*factinv[i]) % mod
	b2 = (b2 + s*i*i*factinv[i]) % mod

a = t*factinv[n]%mod
print((s - 2*t*a + a*a*fact[n])*factinv[n]%mod)