Yukicoder No.2530 Yellow Cards 別解(Bonus 解法)

yukicoder.me

今日 yukicoder contest 411 オムニバスコンテスト - yukicoder というコンテストにて, はちじさん Writer の問題 Yellow Cards が公開されました. 自分は Tester をさせていただきましたが, そこで別解を見つけたので書きたいと思います. この方法では, $N \lt 998244353, K \le 10^{18}$ という制約で解くことが可能であり, 想定解法より高速です.

あらすじ

想定解と少し違う角度で問題を眺めることで, 問題の答えは $K$ について $3$ 次の線形漸化式で表すことができます.

$3 \times 3$ 行列の $K$ 乗を繰り返し二乗法で求め, $F = 3$ として時間計算量 $O(F^{3} \log K)$ などで解くことができます.

解法

次の見方で考えます：イエローカードが $2$ 回目に出されたとき, 仮想的に変数「退場人数」に $1$ を足すだけでその人は退場させず, その人のイエローカードの回数を $0$ にリセットする.

とくに, 退場人数が増える瞬間は「イエローカードの状態が $1 \to 0$ に行ったとき」だけです.

ある $1$ 人についての答えへの寄与を求めたいと考えます. 他の人の寄与も対称性から同じであるため, ある $1$ 人の寄与が求まれば, それを $n$ 倍することで全員の寄与の合計がわかります. ある $1$ 人を A さんとします.

現在の「状態」を, A さんのイエローカードの回数を $2$ で割った余りとします. A さんにイエローカードが出されたとき (確率 $1/n$ ), 状態は $t$ から $t+1 \pmod 2$ に変化します. 他の人にイエローカードが出されたとき (確率 $(n-1)/n$ ), 状態は $t$ から $t$ に変化します.

よって, ${dp}_{i, t}$ を「 $i$ 回目のイベント直後の状態が $t$ である場合の数」とすると, 行列を用いて

$\begin{pmatrix} dp_{i+1, 1}\\ dp_{i+1, 0} \end{pmatrix} = \frac{1}{n} \begin{pmatrix} n-1 & 1\\ 1 & n-1 \end{pmatrix} \begin{pmatrix} dp_{i, 1}\\ dp_{i, 0} \end{pmatrix}$

と表されます. ここで, ケーリー・ハミルトンの定理によって次が分かります：

$\displaystyle dp_{i+2, 1} - \frac{2n-2}{n} dp_{i+1, 1} + \frac{n(n-2)}{n^{2}} dp_{i, 1} = 0$

. いま $dp_{0, 1} = 0, dp_{1, 1} = 1/n$ ということで, $dp_{i,1}$ の母関数は

$\displaystyle \frac{x/n}{1-\frac{2n-2}{n}x+\frac{n(n-2)}{n^2}x^2} = \frac{x}{n(1-x)(1-\frac{n-2}{n}x)}$

となります.

退場人数が増える瞬間は「イエローカードの状態が $1 \to 0$ に行ったとき」であることを思い出します. $i = 1, 2, 3, \cdots, k$ で状態が $1 \to 0$ に行くチャンスがあり, その確率は, $dp_{i-1,1}$ から確率 $1/n$ で起きます. よって, $i$ 回目のイベントで A さんが退場回数に寄与する分の母関数は

$\displaystyle \frac{x^2}{n^2(1-x)(1-\frac{n-2}{n}x)}$

となります. したがって, Aさんの合計の退場回数の期待値は, $i = 1, 2, \cdots, k$ での退場回数への寄与の累積和であることから, 母関数で表すと

$\displaystyle \frac{x^2}{n^2(1-x)^2(1-\frac{n-2}{n}x)}$

となります. よって, 前半に述べたようにこれを $n$ 倍して, 全体での合計退場回数の期待値は

$\displaystyle \frac{x^2}{n(1-x)^2(1-\frac{n-2}{n}x)}$

となります. これは線形漸化式を表し, 行列累乗 (Bostan-Mori などでもよい) で計算できます. 計算量は $F=3$ として $O(F^3 \log K)$ 時間です.

具体的に, 数列 $a$ を $a_0 = 0, a_1 = 0, a_2 = 1/n$ として,

$\displaystyle \begin{pmatrix} a_{K+2}\\ a_{K+1}\\ a_{K} \end{pmatrix} = \begin{pmatrix} \frac{3n-2}{n} & -\frac{3n-4}{n} & \frac{n-2}{n}\\ 1 & 0 & 0\\ 0 & 1 & 0 \end{pmatrix}^K \begin{pmatrix} a_{2}\\ a_{1}\\ a_{0} \end{pmatrix}$