ユーザ解説相当 (2023/08/27 更新)
未履修盛りだくさんで大変だったので、メモを兼ねて $1$ つずつ説明します。というか私はこの手の典型を初めて利用します。
お約束
$0^0=1$ です。 $0^{\underline{0}}=1$ です。形式的冪級数の除算を先頭 $n$ 項について計算量 $O(n\log n)$ で計算できることにします。${}\bmod p$ の逆元の計算量は全体に比べて小さいので、無視します。
問題
yukicoder Advent Calendar Contest 2022 S (=12/18) Paint and Fill
yukicoder No.2166
解説
立式
答えを $A(N,K)$ とします。
$$A(N,K)=\frac{K!}{2^N}[x^K]\left\lparen 2+4x+x^2\right\rparen ^N$$
まずこの式は何ですかというところからです。
マスを赤く塗った個数を $x$ を掛ける回数に対応させながら $N$ 列を走査するとこの式になります。
$1$ 列分の寄与を考えると、
- どちらも塗らない場合が $1$ 通り( $1$ )
- どちらか一方を塗る場合が $2$ 通り( $+2x$ )
- 両方塗る場合が $1$ 通り、ただしカードの並べ方が半減( $+\frac{1}{2}x^2$ )
よって、カードの並べ替え方 $K!$ と合わせて
$$A(N,K)=K![x^K]{\left\lparen 1+2x+\frac{1}{2}x^2\right\rparen }^N$$
が得られ、 $\frac{1}{2}$ を外に出すと最初の式が立ちます。
$K!$ が掛かっているので、 $K\geq 998\ 244\ 353$ の場合の出力は $0$ です。よって、実質 $K\lt 998\ 244\ 353$ を制約にできます。
p-recursive equation の立式
$$A(n,k+1)=(2n-2k)A(n,k)+\frac{k(2n-k+1)}{2}A(n,k-1)$$
これを示します。
ところで、関数 ${\left\lparen f(x)\right\rparen}^n$ を $x$ で微分すると $n{\left\lparen f(x)\right\rparen }^{n-1}f^{\prime}(x)$ です。形式的冪級数でも同様に計算できることが知られています。
ここで $f_n(x)={\left\lparen 1+2x+\frac{1}{2}x^2\right\rparen }^n$ とすると、
$$\begin{aligned}{\left\lparen f_n(x) \right\rparen }^{\prime}&=n{\left\lparen 1+2x+\frac{1}{2}x^2\right\rparen }^{n-1}(2+x)\newline \left\lparen 1+2x+\frac{1}{2}x^2\right\rparen {\left\lparen f_n(x) \right\rparen }^{\prime} &= n{\left\lparen 1+2x+\frac{1}{2}x^2\right\rparen }^n(2+x)\newline \left\lparen 1+2x+\frac{1}{2}x^2\right\rparen {\left\lparen f_n(x) \right\rparen }^{\prime} &= n(2+x)f_n(x) \end{aligned}$$
これを $x^k$ の係数についての方程式にすると、
$$\begin{aligned}[x^k]{\left\lparen f_n(x) \right\rparen }^{\prime}+2[x^{k-1}]{\left\lparen f_n(x) \right\rparen }^{\prime}+\frac{1}{2}[x^{k-2}]{\left\lparen f_n(x) \right\rparen }^{\prime} &= 2n[x^k]f_n(x)+[x^{k-1}]f_n(x)\newline (k+1)\frac{A(n,k+1)}{(k+1)!}+2k\frac{A(n,k)}{k!}+\frac{1}{2}(k-1)\frac{A(n,k-1)}{(k-1)!} &= 2n\frac{A(n,k)}{k!}+n\frac{A(n,k-1)}{(k-1)!}\newline A(n,k+1)+2kA(n,k)+\frac{1}{2}k(k-1)A(n,k-1) &= 2nA(n,k)+nkA(n,k-1) \end{aligned}$$
よって
$$A(n,k+1)=(2n-2k)A(n,k)+\frac{k(2n-k+1)}{2}A(n,k-1)$$
が得られ、つまり $k$ が増加する方向の漸化式を手に入れることができました。
$A(n,-1)=0$ と定義すると、漸化式は $k=0$ の場合も成り立ちます。
難しめの行列累乗
漸化式を行列にします。
$$M_{n,k}=\begin{pmatrix}2n-2k&\frac{k(2n-k+1)}{2}\newline 1&0\end{pmatrix}$$
とすれば、
$$\begin{aligned}\begin{pmatrix}A(n,k+1)\newline A(n,k)\end{pmatrix}&=M_{n,k}\begin{pmatrix}A(n,k)\newline A(n,k-1)\end{pmatrix}\newline \begin{pmatrix}A(n,0)\newline A(n,-1)\end{pmatrix}&=\begin{pmatrix}1\newline 0\end{pmatrix}\end{aligned}$$
ですから、これを累乗すればよいですね、と思いきや、行列に $k$ が入っているので累乗ではありません。正確には、
$$\begin{pmatrix}A(n,k)\newline A(n,k-1)\end{pmatrix}=M_{n,k-1}M_{n,k-2}\ldots M_{n,2}M_{n,1}M_{n,0}\begin{pmatrix}1\newline 0\end{pmatrix}$$
と表されます。
(制約2)多項式の評価点シフト
さきほどの計算を高速にするためのポイントは、多項式を多数の点で同時に評価することです。そのために、次の操作をあらかじめ実現します。
$n-1$ 次以下の多項式 $f$ について、 $f(0),f(1),\ldots ,f(n-1)$ と $t$ と $m$ が与えられたとき、 $f(t),f(t+1),\ldots ,f(t+m-1)$ を計算量 $O((n+m)\log (n+m))$ で計算する。
下降冪を利用します。下降冪は
$$x^{\underline{n}}=x(x-1)\cdots (x-n+1)$$
のように表されます。これを用いた形式的冪級数
$$f(x)=\sum_{k=0}^{\infty} a_kx^{\underline{k}}=a_0+a_1x+a_2x(x-1)+a_3x(x-1)(x-2)+\cdots$$
では、非負整数で評価した値が
$$\begin{aligned}f(0)&=a_00^{\underline{0}}\newline f(1)&=a_01^{\underline{0}}+a_11^{\underline{1}}\newline f(2)&=a_02^{\underline{0}}+a_12^{\underline{1}}+a_22^{\underline{2}}\newline f(3)&=a_03^{\underline{0}}+a_13^{\underline{1}}+a_23^{\underline{2}}+a_33^{\underline{3}}\newline\vdots\end{aligned}$$
のように先頭の有限項のみから定まり、係数列と評価値の間には
$$\sum_{k=0}^{\infty}\frac{f(k)x^k}{k!}=\left\lparen\sum_{k=0}^{\infty}a_kx^k\right\rparen\left\lparen\sum_{k=0}^{\infty}\frac{x^k}{k!}\right\rparen$$
の関係が成り立ちます。(計算に使うのは通常の形式的冪級数です。)また、 $\sum_{k=0}^{n}{(-1)}^k{}_n\text{C}_k=0^n$ より、
$$\left\lparen\sum_{k=0}^{\infty}\frac{f(k)x^k}{k!}\right\rparen\left\lparen\sum_{k=0}^{\infty}\frac{x^k}{k!}{(-1)}^k\right\rparen=\sum_{k=0}^{\infty}a_kx^k$$
として逆の変換ができます。さらに、二項定理のような関係
$$(x+t)^{\underline{n}} = \sum_{k=0}^n {}_{n}\text{C}_k x^{\underline{n-k}} t^{\underline{k}}$$
が成り立ちます(*)。最大次数を $N$ として、 $f(x+t)=\sum_{k=0}^N b_kx^{\underline{k}}$ とすると、これは次のように計算できます。
$$\sum_{k=0}^{N}{(N-k)}!b_{N-k}x^k=\left\lparen\sum_{k=0}^{N}{(N-k)}!a_{N-k}x^k\right\rparen\left\lparen\sum_{k=0}^{N}\frac{t^{\underline{k}}}{k!}x^k\right\rparen\bmod x^{N+1}$$
以上を用いると、評価値から下降冪の級数へ変換して $x$ を平行移動し、評価値に戻すことができます。この手順で高速な評価点のシフトを実現できました。
Library Checker : https://judge.yosupo.jp/problem/shift_of_sampling_points_of_polynomial
Library Checker 平行移動 : https://judge.yosupo.jp/problem/polynomial_taylor_shift
(*) これは Library Checker の Forum に書いてありました。それ以上の調査はしていません。
(制約2)行列積の高速計算
答えの計算式は
$$\begin{pmatrix}A(n,k)\newline A(n,k-1)\end{pmatrix}=M_{n,k-1}M_{n,k-2}\ldots M_{n,2}M_{n,1}M_{n,0}\begin{pmatrix}1\newline 0\end{pmatrix}$$
でした。
このような行列積を処理するためには、まずは行列に含まれている $k$ を不定元 $X$ で解釈し、行列の要素が多項式であると解釈します。このように解釈した行列を $M_{n,X}$ と表します。
まずはいくつかの行列をあらかじめ掛けておきます。すなわち、適当な整数 $A$ を取って
$$L_X=M_{n,X+A-1}\cdots M_{n,X+1}M_{n,X}$$
とすれば、これもまた各要素が多項式であるような行列になります。
すると、行列累乗のようなものは次のように表せます。
$$M_{n,k-1}\cdots M_{n,BA+1}L_{(B-1)A}\cdots L_{2A}L_{A}L_{0}$$
すなわち、 $L_X$ を $X=0,A,2A,\ldots ,(B-1)A$ で評価した行列と高々 $A$ 個の行列の積に変形します。 $A,B\approx \sqrt{K}$ とすることで、行列の計算を $O(\sqrt{K})$ 回にして行列積を計算できました。
問題は、 $X=0,A,2A,\ldots ,(B-1)A$ のときの $L_X$ を、高速に求めることです。
$M_{n,X}$ の要素の最大次数に $1$ を足したものを $d$ とします。
まず $A$ , $B$ を $dA\leq B$ となるように $2$ 冪にとります。ここで $A=2a$ , $B=2b$ , $L^{\prime}_X=M_{n,X+a-1}\cdots M_{n,X+1}M_{n,X}$ とすると、 $L^{\prime}_X$ を $X=0,A,\ldots ,A(b-1)$ で評価した行列から、数回の評価点シフトを用いて、 $L_X$ を $X=0,A,\ldots ,A(B-1)$ で評価した行列を計算できます。この計算は $O(B\log B)$ 時間になります。ここでさらに $a=2a^{\prime}$ , $b=2b^{\prime}$ などとすると、 $A$ 側と $B$ 側を同時に倍々するようなアルゴリズムが得られます。
なぜこのような計算順にするかというと、評価点シフトには多項式の次数よりも多い点で評価した値が必要だからです。 $a$ に対して評価点の個数 $b$ が小さくなると、 $L_X$ の次数が増加して正当性が失われます。今回の例では、初期状態でも行列の多項式は $2$ 次で、 $a$ 個の行列が掛け合わせられているとすれば、その意味で $2a+1\leq b$ が必要です。一方、 $b$ が大きくなりすぎると、評価点シフトの問題サイズが増加して計算量が悪化します。
以上から、制約 $2$ の問題は計算量 $O(T\sqrt{K}\log K)$ で解けます。立式のときに説明した通り、 $K\lt 998\ 244\ 353$ としてよいので間に合います。
(制約2)定数倍高速化
$A$ 側と $B$ 側を同時に倍にするとき、 $4$ つの評価値列をそれぞれ $3$ 回シフトする必要があります。このときの計算は共通する部分があり、省略できます。
評価点シフトの $n=m$ の場合のみ考えます。畳み込みが $3$ 回現れるので、そのままでは $9\times 4\times 3=108$ 回にわたって長さ $2n$ の NTT を計算することになります。そのうち、本当に毎回計算しないといけないのは
- 二項定理の畳み込みの最後の NTT
- 下降冪の多項式から評価値に変換する部分の NTT 2 回
の合計 $3$ 回です。全体で必要になる NTT の回数は $2+1\times 3+3\times 4+3\times 4\times 3=54$ 回になります。
(制約1)多項式の多点評価
$n$ 次の多項式 $P(x)$ を、 $m$ 個の点 $x=a_0,a_1,\ldots ,a_{n-1}$ で評価した値 $P(a_0),P(a_1),\ldots ,P(a_{n-1})$ は計算量 $O((n+m)(\log (n+m))^2)$ で計算できる。
すなわち $P(x)\bmod (x-a_i)$ を求めればよいので、 $P(x)\bmod\prod(x-a_i)$ から始めて多項式除算を利用して分解します。
$(x-a_i)$ を二分木の葉とし、二分木の各ノードはそれ以下の葉の総積をもつとします。まず根について、 $P(x)\bmod\prod(x-a_i)$ を計算します。以下、各ノードでは $P(x)\bmod{}$ (左の子) と $P(x)\bmod{}$ (右の子) を計算します。すると計算時の多項式の次数を抑えながら、 $P(x)\bmod(x-a_i)$ をすべて得られます。
多項式除算を計算量 $O(n\log n)$ で行うには、係数を逆順にした多項式に対して形式的冪級数の除算を計算すればよいです。
Library Checker : https://judge.yosupo.jp/problem/multipoint_evaluation
Library Checker 多項式除算 : https://judge.yosupo.jp/problem/division_of_polynomials
Library Checker FPS の乗法逆元 : https://judge.yosupo.jp/problem/inv_of_formal_power_series
(制約1)行列累積積に対する多点評価
$F_{X,k}=M_{X,k-1}\cdots M_{X,0}$ として、多数の $k=K,X=N$ に対して評価する。
制約 $2$ では $k$ を $X$ で置き換えましたが、今度は $n$ を $X$ で置き換えています。
解くべき問題は多項式の多点評価とは異なりますが、似たアルゴリズムで処理できます。
まず葉では、行列 $M_{X,k}$ または評価点 $(X-N)$ が管理されています。このとき、クエリ $(N,K)$ に対応する葉の左側に、行列 $M_{X,0},M_{X,1},\ldots ,M_{X,K-1}$ がこの順に左から右に向かって並んでいるように配置します。また、各ノードでは行列と評価点それぞれ総積をとったものが管理されています。
ここで、各ノードについて次の行列を求めます。
このノードがもつ評価点を $\mathcal{K}(X)$ 、このノードよりも真に左にある葉の行列の総積を $\mathcal{F}_{X}$ としたとき、 $\mathcal{F}_{X}\bmod \mathcal{K}(X)$ 。ただし、葉がノードの真に左にあるとは、ノードの子孫のうちで最も左の葉と比べて、左にある葉のこと。また、行列 $\bmod$ 多項式とは、行列の各要素で多項式剰余を取ったもの。
これは次のようにトップダウンに計算でき、多項式の多点評価と同様に効率的です。
$\text{visit}(p)\text{ :}$
$\hspace{20px}p.\text{left}.\mathcal{F}_{X}\bmod p.\text{left}.\mathcal{K}(X)=p.\mathcal{F}_{X}\bmod p.\text{left}.\mathcal{K}(X)$
$\hspace{20px}p.\text{right}.\mathcal{F}_{X}\bmod p.\text{right}.\mathcal{K}(X)=\left\lparen p.\text{left}.\textcolor{blue}{F_{X}}\right\rparen\left\lparen p.\mathcal{F}_{X}\right\rparen\bmod p.\text{right}.\mathcal{K}(X)$
$\hspace{20px}\text{visit}(p.\text{left})$
$\hspace{20px}\text{visit}(p.\text{right})$
これにより、制約 $1$ の問題は計算量 $O((T+K){\left\lparen\log (T+K)\right\rparen }^2)$ で解けます。
(制約1)定数倍高速化
行列の各要素を $\mathcal{K}(X)$ で割った剰余を計算しますが、 $\mathcal{K}(X)$ での除算に必要なセットアップは共通化できます。そうすると、わずかに高速化しました。
実装例
定数倍高速化なし
定数倍高速化あり