ECR78 F Cards 解法 - でも今日はSRMあるから

問題

整数n, m, kが与えられる。1/m の確率で成功する事象を独立にn回行う。成功した回数をxとしてpow(x, k)の期待値を MOD 998244353 で求めよ。

n, m < 998244353
k <= 5000

解法

i 回目の事象の確率変数を $x_i$ として、成功ならば $x_i = 1$ 、失敗なら $x_i = 0$ 。問題の期待値は $E[ \left( \sum_i x_i \right) ^k ]$ 。
中を展開すると $n^k$ 個の項の和になる。期待値の線型性より、 $E[ \prod_i x_{a_i} ]$ の和となる。a は解説で言うところの k 要素のtupleで、各要素は1以上n以下。
このtupleの確率変数がすべて1になる確率は 1/m^(a_i の種類数) 。サイズが小さくなったので動的計画法

dp[i][j] := 要素に j 種類含む、長さ i のタプルの個数
dp[0][0] = 1; // 初期値
dp[i+1][j] += dp[i][j] * j; // 既に含まれている要素
dp[i+1][j+1] += dp[i][j] * (n-j); // 新しい要素