rickytheta's blog

競プロ

P-recursive sequenceの第N項の平方分割計算アルゴリズムの理解ログ

これ↓ を理解するまでのログです。

https://nyaannyaan.github.io/library/fps/find-p-recursive.hpp.html

問題定義

扱う数は素数 p を法とした剰余環 Z/pZ とし、以降は単に整数・数と呼びます。

求めたい数列を \{a_i\}_i とし、最終的には a_N を求めます。

数列 \{a_i\}_i は、整数を入力とする高々 d 次の多項式 n+1 個の列 \{p_j(x)\}_{j=0,\dots,n} を用いて、以下の関係が成り立ちます。

 \displaystyle \sum_{j=0}^{n} a_{i+j}p_j(i)=a_{i+n}p_{n}(i)+\dots+a_{i+1}p_1(i)+a_ip_0(i)=0

行列積としての表現

線形漸化式における行列累乗のような考えで、この関係式を行列積として表現し直します。

まず関係式を少し整理します:

 \displaystyle \begin{align*}a_{i+1}p_{n}(i+1-n)&=-\sum_{j=0}^{n-1}a_{i+1-n+j}p_j(i+1-n) \\ a_{i+1}&=\frac{\sum_{j=0}^{n-1}a_{i+1-n+j}p_j(i+1-n)}{-p_n(i+1-n)}\end{align*}

よって、 v_{i} = (a_{i},a_{i-1}, \dots, a_{i-n+1})^T から v_{i+1} = (a_{i+1},a_{i},\dots,a_{i-n})^T に変換するには、 n\times n行列

 \displaystyle \begin{align*}M(x)&=\begin{pmatrix}\{p_{n-1-j}(x+1-n)\}_{j=0,\dots,n-2}&p_0(x+1-n)\\-p_n(x+1-n)\ I_{n-1}&\mathbf0\end{pmatrix}\\&=\begin{pmatrix}p_{n-1}(x+1-n)&p_{n-2}(x+1-n)&\cdots&p_{1}(x+1-n)&p_0(x+1-n)\\-p_n(x+1-n)&0&\cdots&0&0\\0&-p_n(x+1-n)&\cdots&0&0\\\vdots&\vdots&\ddots&\vdots&\vdots\\0&0&\cdots&-p_n(x+1-n)&0\end{pmatrix}\end{align*}

を用いて、 v_{i+1}=\displaystyle\frac{M(i)}{-p_n(i+1-n)}v_{i} と書けます。

求めたいものは a_N、つまり v_N の最初の要素なので、 v_N を求めようと思うと、( p'(i):=-p_n(i+1-n) と定義して)

 \displaystyle v_N=\left(\prod_{i=N-1}^0\frac{M(i)}{-p_n(i+1-n)}\right)v_0=\frac{\left(\prod_{i=N-1}^0M(i)\right)v_0}{\prod_{i=N-1}^0p'(i)}

が計算できれば良いです。(行列積が非可換であるため、総乗は逆順である点に留意します。逆順総乗は記法あってるか分かりませんが雰囲気で……)

なお、ここでは a_i=0 (i\lt 0) としても関係式が成り立つとしていますが、そうでなければ b_i:=a_{i-n+1} と適切にシフトした数列について考え、 v_0 = (a_{n-1},a_{n-2},\dots,a_0) として v_{N-n+1} を求めることを考えれば良いです。

さて、分母については 1\times1 行列 \begin{pmatrix}p'(i)\end{pmatrix} の総乗だと思うと、これは M(i) で n=1 とした場合と構造としては同じになります。よって \prod_{i=N-1}^0 \mathit{Matrix}(i) を計算するアルゴリズムがあれば、これを分母・分子それぞれに1回ずつ適用することで v_N が求まります。

(天下り的な)平方分割の構成

n\times n 行列同士の積の計算量を O(MM(n)) とします。

P([a,b)) := \prod_{i=b-1}^a M(i) = M(b-1)M(b-2)\cdots M(a+1)M(a) と定義します。これは b-a 個の行列の積です。見通しの良さのため半開区間を用います。

ブロックの分割数 B を B:=\lfloor\sqrt{dN}-(d-1)\rfloor とします。この数字は調整可能です。

ブロックの幅 W を W:=\lfloor N/B\rfloor とすると BW\le N なので、 B 個の「 W 個の行列積のブロック」に分けます。

 \displaystyle \begin{align*}&\prod_{i=N-1}^0 M(i)\\=&P([0,N))\\=&P([BW,N))P([(B-1)W,BW))\cdots P([W,2W))P([0,W))\\=&P([BW,N))\prod_{i=B-1}^0P([iW,(i+1)W))\end{align*}

前半にかかっている P([BW,N)) は W>N/B-1 より N-BW\lt B なので愚直に計算すれば O(B\cdot MM(n)) です。よって後半の B 個の総乗 \{ P([iW,(i+1)W)) \}_{i=0,\dots,B-1} を高速に求められれば良いです。

ここまでの図はこんな感じです(説明用の例: N=22, B=5):

この各 P([iW,(i+1)W)) を、以下の図のようなイメージのダブリングによってまとめて計算します(説明用の例: N=22, B=5, d=1):

このダブリングをうまく動かすには、これらの値(行列積)がある多項式によって表現されていることを意識する必要があります。

行列積の多項式表現

乗算される行列の多項式表現を再掲します:

 \displaystyle \begin{align*}M(x)&=\begin{pmatrix}\{p_{n-1-j}(x+1-n)\}_{j=0,\dots,n-2}&p_0(x+1-n)\\-p_n(x+1-n)\ I_{n-1}&\mathbf0\end{pmatrix}\\&=\begin{pmatrix}p_{n-1}(x+1-n)&p_{n-2}(x+1-n)&\cdots&p_{1}(x+1-n)&p_0(x+1-n)\\-p_n(x+1-n)&0&\cdots&0&0\\0&-p_n(x+1-n)&\cdots&0&0\\\vdots&\vdots&\ddots&\vdots&\vdots\\0&0&\cdots&-p_n(x+1-n)&0\end{pmatrix}\end{align*}

p_j(x) は高々 d 次の多項式なので、 p_j(x+1-n) も高々 d 次の多項式です。

例えば M(0) は M(x) に x=0 を代入した時の行列です。

では M(4) について考えると、これも x=4 とすれば良いですが、ここであえて x_4:=x-4 とシフト(変数変換)した行列 M_4(x_4)=M(x+4) を考えます:

 \displaystyle \begin{align*}M_4(x_4)&=\begin{pmatrix}\{p_{n-1-j}(x_4+5-n)\}_{j=0,\dots,n-2}&p_0(x_4+5-n)\\-p_n(x_4+5-n)\ I_{n-1}&\mathbf0\end{pmatrix}\\&=\begin{pmatrix}p_{n-1}(x_4+5-n)&p_{n-2}(x_4+5-n)&\cdots&p_{1}(x_4+5-n)&p_0(x_4+5-n)\\-p_n(x_4+5-n)&0&\cdots&0&0\\0&-p_n(x_4+5-n)&\cdots&0&0\\\vdots&\vdots&\ddots&\vdots&\vdots\\0&0&\cdots&-p_n(x_4+5-n)&0\end{pmatrix}\end{align*}

すると、 M_4(x_4) の各要素も高々 d 次の多項式で、また M(4)=M_4(0) です。

このようにシフトされた行列表現 M_i(x):=M(x+i),M_i(0)=M(i) を考えると、「ある多項式行列の積の x=0 の標本点」と表現し直すことが可能です:

 \displaystyle P([a,b))=\prod_{i=b-1}^a M(i)=\left(\prod_{i=b-1}^a M_i(x)\right)(0)

多項式としての行列積にも表記を与えておきます:

 \displaystyle \begin{align*}P([a,b);x)&:=\prod_{i=b-1}^a M_i(x)\\P([a,b);0)&=P([a,b))\end{align*}

シフトされた行列 M_i(x) の各要素は高々 d 次の多項式だったので、その積である行列 P([a,b);x) はの各要素は高々 d(b-a) 次の多項式になります。

(ここまで次数を執拗に確認していますが、(上界の)次数のみが重要で、実際にどのようなシフトをしてどう展開され係数がいくつになるか、ということには興味が無い点に留意してください。)

また M(x) から M_i(x)=M(x+i) とシフト出来るように、 M_s(x) から M_t(x)=M_s(x+(t-s)) とシフトすることも可能です。これを利用すると、 P([a,b);x) についてもシフトが可能です:

 \displaystyle P([a,b);x+s)=\prod_{i=b-1}^a M_i(x+s)=\prod_{i=b+s-1}^{a+s} M_i(x) = P([a+s,b+s);x)

これを考えると、例えば P([4,8))=P([4,8);0)=P([0,4);4) となります。

ダブリングの構成

さて、行列積の多項式表現と、そのシフトを考えながら図の例に戻してみると、要はこういった構造になっています:

つまり、各ステップで P([0,w); x) という多項式を考え、その各点の標本点(行列)を管理する、という形になっています。 P([0,w); x) の各要素は高々 dw 次の多項式なので、 dw+1 個の標本点を管理すれば、多項式としての情報量を損ないません(ref: ラグランジュ補間)。

よって作られるダブリングアルゴリズムは以下の通りで、最終的に P([0,W); x) の x=0,W,\dots,(B-1)W での標本点が分かれば良いです。

  • 幅 w の行列積 P([0,w); x) の x=0,W,\dots,dwW での標本点 P([0,w);iW) が求まっていると仮定する。
    • 標本点の数は dw+1 個。
    • w=1 のケースは愚直に d+1 個の行列を評価をすることで計算可能。
  • 幅を1増やすには:
    • i\le dw の場合: 各標本点 P([0,w);iW) に単純に行列を掛ける。
      • P([0,w+1);iW)=M(iW+w)P([0,w);iW)
      • dw+1 回の行列評価・行列乗算
    • dw+1\le i\le dw+d の場合: 単純に行列を掛ける。
      • P([0,w+1);iW)=P([iW,iW+w+1))
      • d(w+1) 回の行列評価、 dw 回の行列乗算
    • (動作例: 次数 d=2。w=1 (標本点3個) から w=2 (標本点5個) を w+=1 で作る)

  • 幅を2倍に増やすには:
    • i\le dw の場合: 各標本点を w だけシフトした値を考える。
      • P([0,2w);iW)=P([0,w);iW+w)P([0,w);iW)
      • \{P([0,w);iW+w)\}_{i=0,\dots,dw} は \{P([0,w);iW)\}_{i=0,\dots,dw} に関して、多項式 P([0,w);x) から多項式 P([0,w);x+w) へとシフトした値を考えれば良い。
      • これは多項式の標本点シフトとして知られており、 O(dw\log dw) 回の演算で可能。
      • 標本点シフトに O(n^2 dw\log dw) (行列の各要素の多項式に注目しているので、行列としての演算は発生しない)
      • シフト後の行列を乗算するのに dw+1 回の行列乗算
    • dw+1\le i\le2dw の場合: +(dw+1)W と +(dw+1)W+w だけシフトした値を考える。
      • P([0,2w);iW)=P([0,w);iW+w) P([0,w);iW)
      • j=i-(dw+1) とすると 0\le j\le dw-1
        • P([0,w); iW)=P([0,w);jW+(dw+1)W)
        • P([0,w);iW+w)=P([0,w);jW+(dW+1)W+w)
      • i\le dw と同様に多項式の標本点シフトを2回行うことで計算可能。
      • 標本点シフトに O(n^2 dw\log dw)
      • シフト後の行列を乗算するのに dw 回の行列乗算
    • (動作例: 次数 d=2。w=1 (標本点3個) から w=2 (標本点5個) を w×=2 で作る)

あとは実装するだけです。

多項式の標本点シフト

標本点シフトは以下の問題を解くアルゴリズムです:

  • 高々 d 次の多項式 f(x) の標本点 f(0),f(1),\dots,f(d) が与えられる
  • 整数 c が与えられるので、 f(c),f(c+1),\dots,f(c+d) を求めよ

行間埋めは以下が非常に参考になります:

https://mathlog.info/articles/3351

最終的に3回のFFT畳み込みによって O(d\log d) で高速に計算が可能なことが分かります。

本問題は library-checker-problems (yosupo judge) に収録されているので、具体的な実装はその想定解ファイルを読むのが一番簡便かと思われます。

https://github.com/yosupo06/library-checker-problems/blob/master/math/shift_of_sampling_points_of_polynomial/sol/correct.cpp#L196

今回の問題では f(0),f(W),\dots,f(iW),\dots,f(dwW) という形の標本点になっているので、そのまま使うことは出来ません。

g(x):=f(xW^{-1}) とすると、 f(iW)=g(i) となります。

g(i) として管理をした場合、標本点シフトは f(iW+c)=g(i+cW^{-1}) と考えられ、結局 g(i) に関して cW^{-1} のシフトを実行すれば良いことになります。

残り

なお、「最終的に必要なブロックの数は B」「幅 W の行列積の多項式の次数は dW で、必要な標本点は dW+1」ということを考えると、ダブリングの結果が余ったり足りなかったりします。

  • 余るケース dW+1> B
    • 何も考えずに、最終的に捨てれば良いです。
  • 足りないケース dW+1\lt B
    • k 回の標本点シフトによって標本点の数を k(dW+1) にすることが出来ます。

W:=\lfloor N/B\rfloor だったことを思い出すと、この条件は以下と等価です:

  • 余るケース
    • d\lfloor N/B\rfloor +1\ge d(N/B-1)+1>B
    • B^2+(d-1)B-dN\lt 0
    • B\lt \frac{-(d-1)+\sqrt{(d-1)^2+4dN}}{2} (が成り立つならば余る)
  • 足りないケース
    • d\lfloor N/B\rfloor +1\le d(N/B)+1 \lt B
    • B^2-B-dN \gt 0
    • B>\frac{1+\sqrt{1+4dN}}{2} (が成り立つならば足りない)

特に B\le\sqrt{dN} - (d-1) の時、

 \displaystyle B \le \sqrt{dN}-(d-1) = \frac{-(d-1)+2\sqrt{dN}-(d-1)}{2} = \frac{-(d-1)+\sqrt{(d-1)^2-4(d-1)\sqrt{dN}+4dN}}{2} \le \frac{-(d-1)+\sqrt{(d-1)^2+4dN}}{2}

が成り立つので、 B=\lfloor\sqrt{dN}-(d-1)\rfloor と設定することで、足りないケースは回避することが出来ます。

(もちろん実装時には不等式 d\lfloor N/B\rfloor + 1 \ge B を満たす最大の B を全探索/二分探索で求めた方が効率的です。ここではオーダー把握のため簡易な形式を採用します。)

計算量

最後に、計算量の見積もりを行います。

ある行列 M(i) を評価するのに必要な計算量は、多項式が一番上の行と対角線一個下にしか無いことに注意すると、 O(n^2+nd) です。

行列の乗算に必要な計算量を O(MM(n)) と置いていました。ここではシンプルに乗算するとして O(n^3) ということにします。

\{ P([0,w);iW) \}_{i=0,\dots,dw} を求める計算量を T(w) と置きます。

T(1)=O(d(n^2+nd))=O(dn^2+d^2n) です。

w が奇数の場合は w-1 に再帰します。

T(w)=T(w-1)+O(dw(n^3+nd))=T(w-1)+O(n^3dw+nd^2w) です。

w が偶数の場合は w/2 に再帰します。

T(w)=T(w/2)+O(n^2dw\log dw+ dwn^3)=T(w/2)+O(n^3dw+n^2dw\log dw)) です。

w が高々2ステップで w/2 になることを考えると T(w)=O(n^3dw+n^2dw\log dw+nd^2w) です。

よって T(W)=O(dW(n^3+nd+n^2\log dW)) です。

B=\lfloor\sqrt{dN}-(d-1)\rfloor とすればこのあとに追加で評価点シフトする必要も無いです。

あとは B 個の P([iW,iW+W)) と高々 W 個の M(BW+i) を乗算することで、 \prod_{i=N-1}^0 M(i) が求まります。

この行列積を求める総計算量は

O(dW(n^3+nd+n^2\log dW)+(B+W)n^3)=O(n^2dW\log dW+n^3(B+dW)+nd^2W)

となります。 n,d \ll W,N とすれば O(n^2dW \log dW) で、 W=\lfloor N/B\rfloor=O(\sqrt{N/d}) を思い出すと O(n^2\sqrt{dN}\log dN) です。

(なお、NTT-friendly でない素数の場合、畳み込みに garner's algorithm 等を用いた任意mod畳み込みを使用する必要があります。garner's algorithm を用いる場合 \bmod\ p での逆数計算が必要になるため、 O(n^2\sqrt{dN}(\log p+\log dN)) になります。)

その他

  • 本アルゴリズムでは「たかだか dw 次の多項式は dw+1 個の標本点を保存すれば良い」という性質を用いて、標本点の数自体も幅 w と一緒にダブリングしていますが、最初から最終的な多項式の次数 dW に対応した dW+1 個の標本点を持ち続けるアルゴリズムも設計可能です。
    • この場合、標本点の数が最初から最後まで O(dW) なので、段数がかかって管理する標本点の総数が O(dW\log\cdot) になります。
    • 標本点の数自体もダブリングする方式では級数を考えることで、標本点の総数が O(dW) になるので、この分の \log が落ちることになります。
  • また、多項式 P([0,W);x) を陽に求め、多点評価によって P([0,W);iW) を求めるアルゴリズムも設計可能です。

後記

  • お久しぶりです
    • また O(\sqrt N) の記事書いてる……
    • 関係ないですがTTPC良かったですね。僕はTTPC upsolveに睡眠時間を削られています。来年はTTPC(という名前のコンテスト)が開催されるのでしょうか……
  • P-recursive がちょうど流行り始めてそうな2020年～2021年頃にぬるっと引退したので存在を知らず、すげ～となった
    • でも N! \bmod p 自体はもっと昔から存在はしていたらしいので、単純に精進が足りなかっただけでもありそう
    • 雑に Berlekamp-Massey に投げる感覚で P-recursive 判定に投げる世界、ヤバい
  • 階乗を o(N) で解けるのだいぶ驚きだったんだけど、 o(\sqrt N) はやばそうだな……という感覚がある
    • 素因数分解が解けてしまうので……
    • と思ったらそれが Strassen’s Factoring algorithm だった
  • FPSの文脈で出がちなイメージだけど、おそらくコレ自体は”FPS”とは関係ない
    • x に値をバリバリ代入するし、やってることはただの多項式演算
    • 係数に意味があるわけでもない(多分、おそらく; 研究的価値はありそう)
    • とはいえ結局数列 \{a_i\}_i を錬成するのにFPS使いがちなのでジャンルとしてはFPSの延長っぽい

参考文献