yukicoder No. 2530: Yellow Cards - SSRS の競プロ日記

$N \not\equiv 0 \pmod {998244353}$ のもとで、 $O(\log N + \log K)$ 時間で解きます。

解法

[1] $\Theta(N+K^2)$ 解法

$dp[i][j]$ を、 $i$ 回のイベントの直後に退場となった選手の人数がちょうど $j$ 人であるような確率とします。

$i$ 回のイベントの直後に退場となった選手の人数がちょうど $j$ 人であるとき、今までの $i$ 回のイエローカードの提示のうち $2j$ 回は選手を退場させるために使われているので、プレイエリアにいる選手のうちイエローカードを $1$ 回提示された選手の人数は $i-2j$ 人となります。よって、以下のような漸化式が得られます。

$dp[i+1][j] += \displaystyle\frac{N-(i-2j)}{N} dp[i][j]$
$dp[i+1][j+1] += \displaystyle\frac{i-2j}{N} dp[i][j]$

また、初期値は $dp[0][0] = 1$ 、答えは $\displaystyle\sum_{i=0}^{\infty} (N+i) \cdot dp[K][i] = N + \displaystyle\sum_{i=0}^\infty i \cdot dp[K][i]$ となります。(無限和の形で書いていますが、実際には加える要素数は $O(K)$ 個となります。)

[2] $\Theta(N+K)$ 解法

$\displaystyle\sum_{j=0}^\infty dp[i][j] = 1$ であることに注意します。 $\displaystyle\sum_{j=0}^\infty j\cdot dp[i][j]$ の値に関心があるので、この値を $dp2[i]$ とおきます。

$dp$ の漸化式から、
$\begin{align*} dp2[i+1] &= \sum_{j=0}^\infty \left(j\displaystyle\frac{N-(i-2j)}{N} + (j+1)\displaystyle\frac{i-2j}{N}\right)dp[i][j] \\ &= \sum_{j=0}^\infty \left(j+\displaystyle\frac{i-2j}{N}\right)dp[i][j] \\ &= \displaystyle\frac{N-2}{N}\sum_{j=0}^\infty j\cdot dp[i][j] + \frac{i}{N}\sum_{j=0}^\infty dp[i][j] \\ &= \displaystyle\frac{N-2}{N}dp2[i]+\frac{i}{N} \end{align*}$

この漸化式に従い、 $dp2[i+1]$ を $dp2[i]$ から $O(1)$ 時間で計算できます。

[3] $\Theta(\log N + \log K)$ 解法

$dp2$ の漸化式を解きます。特殊解として $dp2[i] = ai+b$ の形を仮定すると、

$\begin{align*} a(i+1)+b &= \displaystyle\frac{N-2}{N}(ai+b) + \displaystyle\frac{i}{N} \\ \displaystyle\frac{2}{N}(ai+b)+a &= \displaystyle\frac{i}{N} \\ \displaystyle\frac{2a-1}{N}i + \left(a+\displaystyle\frac{2}{N}b\right) &= 0 \end{align*}$

両辺を係数比較することで $a = \displaystyle\frac{1}{2}, b = -\displaystyle\frac{N}{4}$ を得ます。よって、 $dp2[i+1]-\left(\displaystyle\frac{1}{2}(i+1)-\displaystyle\frac{N}{4}\right) = \displaystyle\frac{N-2}{N}\left(dp2[i]-\left(\displaystyle\frac{1}{2}i-\displaystyle\frac{N}{4}\right)\right)$