I - Back to Origin

目次

まとめ：まず、形式的冪級数を用いて問題を整理し、合成関数を求める問題に帰着します。ここで求めたい合成関数 $A(B(x))$ の中身 $B(x)$ が特殊な形をしていることを利用して、高速に合成を求めることができます。 $O(N \log N)$ または $O(N \log^2 N)$ 時間でこの問題に正解することができます。

• 目次
• 考察
  • 導入
  • $B(x)$ の表示
  • $D(x)$ の表示
  • まとめ
• $O(N \log N)$ の方法…想定解(Pypy3 でも通る)
  • 前半戦
  • 後半戦
• $O(N \log^2 N)$ の方法 1…想定解
  • 漸化式・行列
  • 分割統治
  • 定数倍高速化
• $O(N \log^2 N)$ の方法 2…新しい解法（合成を求める）
• $O(N \log^2 N)$ の方法 3…新しい解法
  • 定数倍高速化

考察

導入

問題の答えを $a_1$ とします。操作の 「$-1, +1$ から等確率で選ぶ」の部分を「$-1, +1$ のどちらか一方を選ぶ」に変更したとき、ありえる操作すべてに対するスコアの総和を $a_2$ とします。このとき、 $a_1 = a_2 / 2^N$ であることが分かります。 $1/2^N \bmod{998244353}$ は容易に計算できるため、代わりに $a_2$ を計算することを考えます。

しょぼん君の操作による動きをパスと呼ぶことにします。次のように形式的冪級数を定めます。

• $B(x)$ は、座標 $0$ から出発し $n$ ステップ目で座標 $0$ に初めて戻ってくる長さ $n$ のパスの個数が $x^n$ の係数になるような形式的冪級数
• $D(x)$ は、座標 $0$ から出発して座標 $0$ に一度も戻らないような長さ $n$ のパスの個数が $x^n$ の係数になるような形式的冪級数
• $A(x)$ は、 $A(x) = A_0 + A_1 x + \cdots + A_N x^N$ を満たす多項式
• $F(x)$ は、長さ $n$ のパスすべてに対するスコアの合計が $x^n$ の係数になるような形式的冪級数

求めたいものは $[x^N] F(x)$ です。ここで、「座標 $0$ にちょうど $k$ 回戻ってくる」、すなわちスコア $A_k$ を獲得できるようなパスの個数の母関数は

$B(x)^kD(x)$

と表されるため、 $F(x)$ は次のように表せます。

$F(x)\\ = A_0 D(x) + A_1 B(x) D(x) + A_2 B(x)^2 D(x) + \cdots + A_N B(x)^N D(x)\\ = D(x)A(B(x))$

$B(x)$ の表示

$B(x)$ と $D(x)$ を具体的に表示してみます。 $B(x)$ は「座標 $0$ から出発し、 $0$ に初めて戻って来る」というものでした。もし $1$ 回目で座標 $1$ に行ったとすると、$2$ から $n-1$ 回目では座標 $1$ 以上のみを動き、$n$ 回目で座標 $1 \to 0$ にもどってくる、ということになります。そして、 $1$ 回目で座標 $-1$ に行ったときも同様です。

$2$ から $n-1$ 回目の全 $n-2$ 回において、これは「座標 $0$ からスタートし、座標 $0$ 以上を $n-2$ 回動き、座標 $0$ に戻ってくる」方法の数と同じで、この場合の数はカタラン数と同じです。 $n$ が奇数のときは $0$ 通り、偶数のときは $\frac{1}{(n-2)/2 + 1}\binom{n-2}{(n-2)/2}$ となります。

カタラン数の母関数 $C(x)$ を

$C(x)\\ = \sum_{n\ge 0} \frac{1}{n+1} \binom{2n}{n} x^n\\ = 1 + x + 2x^2 + 5x^3 + 14x^4 + 42x^5 + \cdots$

とおくと、$B(x)$ は

$B(x)\\ = 2x^2 C(x^2)\\ = 2x^2 + 2x^4 + 4x^6 + 10x^8 + 28x^{10} + 84x^{12} + \cdots$

と表されます。ここで、カタラン数の母関数は

$\displaystyle C(x) = \frac{1 - \sqrt{1-4x}}{2x}$

であるので、

$B(x) = 2x^2 C(x^2) = 1 - \sqrt{1-4x^2}$

が分かりました。

$D(x)$ の表示

次に、$D(x)$ について、余事象は「$0$ からスタートし、 $0$ に戻ってくる瞬間が存在する」ですが、 $0$ に初めて戻ってくるのが何回目であるかで場合分けすれば、

$[x^n] D(x) = 2^n - [x^n] B(x) - 2[x^{n-1}] B(x) - \cdots - 2^n [x^0] B(x)$

が得られます。これはとくに、

$\displaystyle [x^n] D(x)\\ = 2^n - [x^n] \frac{B(x)}{1-2x}\\ = [x^n] \frac{1-B(x)}{1-2x}$

であるため、

$\displaystyle D(x) = \frac{\sqrt{1-4x^2}}{1-2x}$

が分かります。

まとめ

以上を整理すると、求めるべきものは

$\displaystyle [x^N] F(x)\\ = [x^N] D(x) A(B(x))\\ = [x^N] \frac{\sqrt{1-4x^2}}{1-2x} A \left( 1- \sqrt{1-4x^2}\right)$

になります。

$O(N \log N)$ の方法

前半戦

今回の場合、合成関数の中身 $1 - \sqrt{1-4x^2}$ が簡単であるため、合成関数を直接「ほどく」イメージで求められます。

$A(1-\sqrt{1-4x^2})$ に対して次のように式変形します。

$\displaystyle A(1 - \sqrt{1-4x^2})\\ = \sum_{i=0}^N A_i (1 - \sqrt{1-4x^2})^i\\ = \sum_{i=0}^N A_i \sum_{k=0}^i \binom{i}{k} (- \sqrt{1-4x^2})^{k}\\ = \sum_{k=0}^N (-\sqrt{1-4x^2})^{k} \sum_{i=k}^N A_i \binom{i}{k} \\ = \sum_{k=0}^N (-\sqrt{1-4x^2})^{k} [x^k] \left(\sum_{i=0}^N A_i (1+x)^i\right)$

ここで、 $\sum_{i=0}^N A_i (1+x)^i = P_0 + P_1 x + \cdots + P_N x^N$ とおきます。そうすると、

$\displaystyle A(1 - \sqrt{1-4x^2}) = \sum_{k=0}^N P_k (-\sqrt{1-4x^2})^{k}$

です。そして、 $P_0, P_1, \cdots, P_N$ は次のように $O(N \log N)$ 時間で求められます：

$I(x) = \sum_{i=0}^N A_i x^i$

とおきます。このとき、求めたい多項式は $I(x+1)$ であり、その各項は Polynomial Taylor Shift というテクニックを用いて求められます。

Polynomial Taylor Shift については、 ABC215-G のユーザー解説 (MMNMM さん) をご覧ください。

後半戦

$\displaystyle A(1 - \sqrt{1-4x^2}) = \sum_{k=0}^N P_k (-\sqrt{1-4x^2})^{k}$

の各係数を求めたいです。 $k$ の偶奇で分け、

$\displaystyle \sum_{k=0}^N P_k (-\sqrt{1-4x^2})^{k}\\ = \sum_{k=0}^{\lfloor N/2 \rfloor} P_{2k} (1-4x^2)^{k} - \sqrt{1-4x^2} \sum_{k=0}^{\lfloor (N-1)/2 \rfloor} P_{2k+1} (1-4x^2)^{k}$

とします。ここで、

$\sum_{k=0}^{\lfloor N/2 \rfloor} P_{2k} (1-4x^2)^{k} = X_0 + X_1 x + \cdots + X_N x^N$

$\sum_{k=0}^{\lfloor (N-1)/2 \rfloor} P_{2k+1} (1-4x^2)^{k} = Y_0 + Y_1 x + \cdots + Y_N x^N$

とおくと、

$\displaystyle A(1 - \sqrt{1-4x^2})\\ = (X_0 + X_1 x + \cdots X_N x^N) - \sqrt{1-4x^2} (Y_0 + Y_1 x + \cdots Y_N x^N)$

となります。いま、 $X_i, Y_i$ の各項をすべて求めることは、同じく Polynomial Taylor Shift によって $O(N \log N)$ 時間で出来ます。

そして、 $\sqrt{1-4x^2} = (1-4x^2)^{1/2}$ は疎な形式的冪級数の pow であるので、$x^N$ までの全ての項が $O(N)$ 時間で計算できます。あとは $\sqrt{1-4x^2}$ と $Y_0 + Y_1x + \cdots + Y_Nx^N$ との畳み込みを行えば、 $A(1 - \sqrt{1-4x^2})$ の各項が計算できました。

疎な形式的冪級数のpowを高速に求めるテクニックについては A problem collection of ODE and differential technique (jqdai0815 さん) の「Exp in Gennady Korotkevich Contest 5 By tourist」の節をご覧ください。

最後に、 $D(x) A(1-\sqrt{1-4x^2})$ が求める答えです。$D(x) = \frac{\sqrt{1-4x^2}}{1-2x}$ は $\sqrt{1-4x^2}$ の累積和を求めるのと同じようにして $O(N)$ 時間で求められます。その後、 $D(x)$ と $\sqrt{1-4x^2}$ との畳み込みを行えばよいです。以上から、合わせて $O(N \log N)$ 時間で解けます。

$O(N \log^2 N)$ の方法

漸化式・行列

$B(x)$ はカタラン数とほぼ同じです。カタラン数の性質を用いることを考えます。カタラン数の母関数 $C(x)$ について、

$x C(x)^2 = C(x) - 1$

が成り立ちます。これを用いると、 $B(x) = 2x^2 C(x^2)$ について

$B(x)^2\\ = 4x^4 C(x^2)^2\\ = 4x^2 (C(x^2) - 1)\\ = 2 B(x) - 4x^2$

が成り立つことが分かります。求めるべきは

$\displaystyle A(B(x)) = \sum_{i=0}^{N} A_i B(x)^i$

です。これから、先ほどの漸化式を用い、分割統治を使ってこの式を高速に求めます。

多項式 $P_k(x), Q_k(x)$ を用いて $B(x)^k = P_k(x) B(x) + Q_k(x)$ とおくと、

$\displaystyle \sum_{i=0}^{N} A_i B(x)^i = \sum_{i=0}^{N} A_i (P_i B(x) + Q_i)$

となります。これを高速に求めることができるでしょうか？$(P_i, Q_i)$ の組から $(P_{i+1}, Q_{i+1})$ の組への遷移は次のようになります。

$\displaystyle B(x)^{i+1}\\ = B(x) B(x)^i\\ = B(x) (P_i B(x) + Q_i)\\ = P_i (2B(x) - 4x^2) + Q_i B(x)\\ = (2P_i + Q_i) B(x) - 4x^2 P_i$

よって、多項式を係数とした行列を用いて

$\begin{pmatrix}P_{i+1}\\Q_{i+1}\end{pmatrix} =\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix} \begin{pmatrix}P_{i}\\Q_{i}\end{pmatrix}$

と表されます。とくに、

$\displaystyle \sum_{i=0}^{N} A_i B(x)^i\\ = \sum_{i=0}^{N} A_i \begin{pmatrix}B(x)\\1\end{pmatrix} \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^i \begin{pmatrix}0\\1\end{pmatrix}\\ = \begin{pmatrix}B(x)\\1\end{pmatrix} \left (\sum_{i=0}^{N} A_i \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^i \right) \begin{pmatrix}0\\1\end{pmatrix}$

であるため、$\displaystyle \sum_{i=0}^{N} A_i \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^i$ を求めればよいです。

分割統治

$\displaystyle f(L, R) = \sum_{i=L}^{R-1} A_i \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{i-L}$

とします。求めるものは $f(0, N+1)$ です。

まず、 $R - L = 1$ の場合は $f(L, R) = \begin{pmatrix}A_L&0\\0&A_L\end{pmatrix}$ です。

そうでない場合、 $\lfloor (L+R)/2 \rfloor = M$ として、$f(L, M), f(M, R)$ が求まったとします。このとき、 $f(L, R)$ が次のように計算できます：

$\displaystyle f(L, R)\\ = \sum_{i=L}^{R-1} A_i \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{i-L}\\ = \sum_{i=L}^{M-1} A_i \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{i-L}+\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{M-L} \sum_{i=M}^{R-1} A_i \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{i-M}\\ = f(L, M) + \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{M-L} f(M, R)$

ここで、 $\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{M-L}$ と $f(M, R)$ との畳み込みで $O(N \log N)$ 時間かかります。

$\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{M-L}$ については同じく分割統治で求められます。

この分割統治は $f(L, R)$ の統治の部分に $O((R-L) \log(R-L))$ 時間かかるので、 $f(0, N+1)$ を求めるのに計算量は $O(N \log^2 N)$ 時間かかります。

これで合計 $O(N \log^2 N)$ 時間でこの問題が解けました。

定数倍高速化

このままでは定数倍が重く、C++ であっても通すのが難しいです。しかし、次のような定数倍改善がいくつか考えられ、これらを使って AC を獲得することができます。（1, 2, 3 番目のどれかを用いるのが簡単です）

• 原点に戻ってくる回数は高々 $\lfloor N/2 \rfloor$ 回であるため、$A$ を $A_0, \cdots, A_{\lfloor N/2 \rfloor + 1}$ しかないことにしてよい。
• $A$ の長さを $2$ べき、すなわちある非負整数 $k$ が存在して $A_{N+1} = \cdots = A_{2^k-1} = 0$ とする。これによって、分割統治の際 $\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{M-L}$ における $M-L$ がすべて $2$ べきになる。$\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{1}, \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{2}, \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{4}, \cdots, \begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{2^i}, \cdots$ は分割統治をせずに合計 $O(N \log N)$ で求めることができる。
• $\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}$ の代わりに $\begin{pmatrix}2&1\\-4x&0\end{pmatrix}$ を用いて計算してよい。後で $x$ を $x^2$ に置き換えれば求めたいものが求まる。
• 行列積を計算する際、FFT を $8$ 回行うと butterfly が(inv含めて)合計 $16$ 回行われるが、butterfly の回数を $8$ 回に削減することが可能である。$A$ の長さを $2$ べきにすることで、$\begin{pmatrix}2&1\\-4x^2&0\end{pmatrix}^{2^i}$ の butterfly が再利用でき、さらに定数倍改善することも可能である

$O(N \log^2 N)$ の方法 2

noshi91 さんが2024年3月16日に投稿した記事 FPS の合成と逆関数、冪乗の係数列挙 Θ(n (log(n))^2) で、合成関数を求める操作が $O(N \log^2 N)$ 時間で出来ることが示されています。そして、実際に合成を $O(N \log^2 N)$ 時間で行うアルゴリズムによって、この問題を AC できます。

$O(N \log^2 N)$ の方法 3

noshi91 さんが2024年3月16日に投稿した記事 FPS の合成と逆関数、冪乗の係数列挙 Θ(n (log(n))^2) と同じ方法で、固定された $k$ に対して $[x^k] a(x) b(x)^i$ をすべての $0 \le i \lt n$ に対して列挙することが $O(n \log^2 n + k\log^2 k)$ 時間で可能です。 $2$ 変数形式的冪級数 $F\llbracket y \rrbracket \llbracket x \rrbracket$ において $\displaystyle [x^k] \frac{a(x)}{1-yb(x)}$ の $y^n$ 次までを求めればよく、bostan-mori 法を用いて求められます。（詳しくは記事をご覧ください）

今回、求めたいものは

$\displaystyle [x^N] D(x) A(B(x))\\ = \sum_{i=0}^N A_i D(x) B(x)^i$

でした。よって、$\displaystyle [x^N] \frac{D(x)}{1-yB(x)}$ を $y^N$ まで求めればこの問題を $O(N \log^2 N)$ 時間で解くことができます。

定数倍高速化

この方法だと定数倍が重く、実行時間制限に間に合わない可能性があります。そこで、次の ad-hoc な定数倍高速化が考えられます。やることは bostan-mori 法の一部分とほとんど同じです。

• この問題の場合、 $B(x)$ は偶関数であるので、 $B(x) = V(x^2)$ となる形式的冪級数 $V$ が存在する。また、 $D(x)$ を $2$ つの偶関数 $D_e(x^2), D_o(x^2)$ に分け、 $D(x) = D_e(x^2) + xD_o(x^2)$ とおく。このとき

$\displaystyle [x^N] \frac{D(x)}{1-yB(x)} \\ = [x^N] \frac{D_e(x^2) + xD_o(x^2)}{1-y V(x^2)}\\ = \begin{cases} [x^N] \frac{D_e(x^2)}{1-y V(x^2)} & (N は偶数)\\ [x^N] \frac{x D_o(x^2)}{1-y V(x^2)} & (N は奇数) \end{cases}\\ = \begin{cases} [x^{N/2}] \frac{D_e(x)}{1-y V(x)} & (N は偶数)\\ [x^{(N-1)/2}] \frac{D_o(x)}{1-y V(x)} & (N は奇数) \end{cases}$

が成り立つ。また、実際には $y$ は $y^{N/2}$ 程度までしか列挙する必要がない。最初にこの処理をしておくと、 $2$ 倍以上の定数倍高速化となる。