ABC473 G - Wipeout

考え方

各カード、$3$ 回以上めくる必要はない。
$1$ 回めくってハズレだった場合、それを食べるチャンスが来るまでめくらず放っておけばよい。
しかも、最後の $1$ 枚はめくらずとも中身がわかるので、ハズレになることはない。
ということで、まず $K\ge 2N$ だった場合は、即 $0$ と答えてしまって問題ない。
以下では $K<2N$ である前提で考える。

各カード、$1$ 回は必ずめくって食べることになる。
ということは、この問題は、めくったが食べられないことが $K-N$ 回ある確率を求めればよい。

そして、めくったことがないカード同士は対等である。
よって、次にめくるべきカードが未出である場合、めくったことがない最も左のものをめくることにする。

さて、0-indexed で数えて左から $i$ 枚目をめくることを考える。
このカードを初めてめくったときに一発で食べることができる確率を求める。
一発で食べられる必要十分条件は、残っている $N-i$ 枚の中で最小のカードが最も左にあることである。
よって、その確率は $1/(N-i)$ である。

これで、$O(N^2)$ でよければ、$i$ 番目までに $j$ 回失敗する確率を考える動的計画法で解ける。
ただし、この問題の制約では、この計算量は間に合わない。

そこで、形式的冪級数を用いて高速化する。
求める値は、以下の式の $x^{K-N}$ の係数である。

$$
P(x)
=
\left(\dfrac{1}{N}+\dfrac{N-1}{N}x\right)
\left(\dfrac{1}{N-1}+\dfrac{N-2}{N-1}x\right)
\dots
\left(\dfrac{1}{2}+\dfrac{1}{2}x\right)
\left(\dfrac{1}{1}+\dfrac{0}{1}x\right)
$$

整数で計算したいので、これを以下のように変形する。

$$
P(x)
=
\{1+(N-1)x\}
\{1+(N-2)x\}
\dots
(1+1x)
(1+0x)
\times \dfrac{1}{N!}
$$

この $N$ 個の積を畳み込みで計算し、$x^{K-N}$ の係数を $N!$ で割れば答えとなる。
畳み込みは $2$ 個ずつ行うが、長いものを何度も扱うと計算量がかさんでしまう。
残っている中で短い方から $2$ つの積をとることを繰り返して計算量を抑えること。
実装的には、全てを queue にいれてから、前 $2$ つの畳み込みを後ろに入れるループでよい。

この高速化をきちんとすれば、計算量は $O(N\log^2 N)$ である。

入力例1での動作

入力を受け取る。

n: 3
k: 4

$K-N=4-3=1$ なので、失敗回数が $1$ 回となる確率を求める。

整数係数側の多項式は $(1+2x)(1+x)(1+0x)=1+3x+2x^2$ となる。

求めるのは $x^1$ の係数なので $3$ である。
これを $N!=3!=6$ で割ると、$3/6=1/2$ となる。

法 $998244353$ で $1/2$ は $499122177$ なので、これを出力する。

注意点

答えは $998244353$ で割った余りを要求されているので、剰余類環の考えに従って処理する。
何か足し算や掛け算をするたびに結果を % 998244353 し、割り算は $998244353-2$ 乗したものを掛ける。
累乗計算を行うため、繰り返し二乗法のアルゴリズムも用意しておくこと。

別解

特になし。