Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 243 F問題 Lottery

AtCoder ABC ABC243 F問題 500 pts

問題

提出解答

問題の概要

1回引くと, $N$ 種類の賞品のうちどれか1種類が1個貰えるくじがある. $i$ 番目の賞品が貰える確率は $\displaystyle \dfrac{W_i}{\sum_{i=1}^N W_i}$ である.

また, 各回は独立である.

このくじを $K$ 回引いたあと, 貰えた賞品がちょうど $M$ 種類である確率を求めよ.

制約

$1 \leq K \leq 50$
$1 \leq M \leq N \leq 50$
$0 \lt W_i$
$\displaystyle 0 \lt \sum_{i=1}^N W_i \lt 998244353$ *1

解法

1回のくじで $i$ 番目の賞品が貰える確率を $\displaystyle p_i=\dfrac{W_i}{\sum_{j=1}^N W_j}$ とする.

求めるべきは

$\displaystyle \sum_{\substack{0 \leq k_i \\ k_1+\dots+k_N=K \\ \#\{1 \leq i \leq N \mid k_i \gt 0 \}=M }} \dbinom{K}{k_1, \dots, k_N} p_1^{k_1} \dots p_N^{k_N}=K! \sum_{\substack{0 \leq k_i \\ k_1+\dots+k_N=K \\ \#\{1 \leq i \leq N \mid k_i \gt 0 \}=M }} \prod_{i=1}^N \dfrac{p_i^{k_i}}{k_i!}$

である. $\#\{1 \leq i \leq N \mid k_i \gt 0 \}=M$ という条件がなければ, 1変数の形式的ベキ級数を用いて比較的簡単に求めることができる.

しかし, この問題ではこの制約により, 1変数で議論することが厳しくなっている. ここで, $X$ を変数としたとき, $X^k$ で状態を表していることが頭に入っていると, もう1つの変数を増やして, 2つの状態を記録すればよいうことに気がつく.

そこで, 2変数の形式的ベキ級数を用いて求めることにする. $X,Y$ を変数とする形式的ベキ級数 $F_i(X,Y)$ を

$\displaystyle F_i(X,Y)=1+Y \sum_{n=1}^\infty \dfrac{p_i^n}{n!} X^n$

と定める. このとき, $X^n Y^r$ 次は $n$ 回引いて $r$ 種類の賞品を獲得した状況を意味する.

このことから, 求めるべき答えは

$\displaystyle K! \left(\left[X^K Y^M\right ] \prod_{i=1}^N F_i (X,Y) \right)$

である.

2変数の形式的ベキ級数の保存方法について,

$\displaystyle P(X,Y)=\sum_{n=0}^\infty P_n(X) Y^n$

とみなして, 1変数の形式的ベキ級数の列として保存すればよい.

また, 2つの2変数形式的ベキ級数 $P(X,Y), Q(X,Y)$ の積について,

$\displaystyle P(X,Y) Q(X,Y)=\sum_{n=0}^\infty \left(\sum_{m=0}^n P_m(X) Q_{n-m}(X) \right) Y^n$

であり, 1変数の場合に帰着できる.

以上では, 形式的ベキ級数での議論を行ったが, 実際には $X$ については $K$ 次まで, $Y$ については $M$ 次まででよく. また, 多項式の積についても愚直に行っても全体の時間計算量 $O(NMK^2)$ で求めることができる.

また, $P_m(X) Q_{n-m}(X)$ の部分を数論変換を利用すると, 時間計算量 $O(NMK \log K)$ になる.

*1: $\pmod{998244353}$ での出力を求められているため