Ex - Walk Editorial by Nyaan
頂点 \(1\) へ向かう辺が存在しない場合
はじめに、頂点 \(1\) に向かう辺が存在しない場合を考えます。つまり \(D_n = 0, A_1 = 0\) です。
- \(\mathrm{dp}_{n, m} : \) (頂点 \(1\) から頂点 \(n\) へ向かう \(m\) 辺 walk の本数)
とすると、 \(\mathrm{dp}_{n,m}\) は次の式で表せます。 (\([\mathrm{cond}]\) は \(\mathrm{cond}\) が真の時 \(1\), 偽の時 \(0\) を取る関数) また、計算したいものは \(\mathrm{dp}_{N, K}\) です。
\[ \begin{aligned} \mathrm{dp}_{0, 0} &= 1 \\ \mathrm{dp}_{n,m} &= \mathrm{dp}_{n,m-1}[A_n = 1] \\ &+ \mathrm{dp}_{n-1,m-1}[B_{n-1} = 1] \\ &+ \mathrm{dp}_{n-2,m-1} [C_{n-2} = 1] \end{aligned} \]
母関数で記述します。\(F_n(x) = \sum_{m} \mathrm{dp}_{n,m} x^m\) とすると次の式を得ます。(\(A_1 = 0\) に注意)
\[ \begin{aligned} F_1(x) &= 1 \\ F_n(x) &= x F_n(x) [A_n=1] \\ &+ x F_{n-1}(x) [B_{n-1} = 1] \\ &+ x F_{n-2}(x) [C_{n-2} = 1] \\ \iff F_n(x) &= \left(\frac{1}{1-x [A_n = 1]}\right) \\ & \times(x F_{n-1}(x) [B_{n-1} = 1]+ x F_{n-2}(x) [C_{n-2} = 1]) \end{aligned} \]
見やすさのため \(P_i(x), Q_i(x), R_i(x)\) を
\[ \begin{aligned} P_n(x) &= \frac{1}{1-x [A_n = 1]} \\ Q_n(x) &= x [B_{n-1} = 1] \\ R_n(x) &= x [C_{n-2} = 1] \end{aligned} \]
と置いて式を書き換えると
\[ \begin{aligned} F_1(x) &= 1 \\ F_n(x) &= P_n(x) \times(Q_n(x) F_{n-1} (x) + R_n(x) F_{n-2}(x)) \end{aligned} \]
という関係式を得ます。
この漸化式を計算すれば \(\mathrm{dp}_{N, K} = \lbrack x^K \rbrack F_N(x)\) を求められます。愚直に漸化式を計算すると \(\Omega(N^2)\) 掛かりますが、 行列で表現すると
\[ \left( \begin{array}{c} F_{n+1} \\ F_{n} \end{array} \right) = \left( \begin{array}{cc} P_n Q_n & P_n R_n \\ 1 & 0 \end{array} \right) \left( \begin{array}{c} F_n \\ F_{n-1} \end{array} \right) \]
\[ \left( \begin{array}{c} F_{N} \\ F_{N-1} \end{array} \right) = \left( \begin{array}{cc} P_N Q_N & P_N R_N \\ 1 & 0 \end{array} \right) \cdots \left( \begin{array}{cc} P_2 Q_2 & P_2 R_2 \\ 1 & 0 \end{array} \right) \left( \begin{array}{c} F_{1} \\ 0 \end{array} \right) \]
という関係式が得られて、行列の各係数は分子・分母が \(1\) 次以下の有理式で表せるので、行列積の部分で二分木風にマージする形で積を計算すると \(\mathrm{O}(N \log^2 N)\) で \(F_N\) を求められます。 (なお、\(\frac{1}{1-x}\) が出て来る行列にあらかじめ全体に \((1 - x)\) を掛けておいて、総積を取った後に \((1-x)\) の冪でまとめて割ることにすると、多項式係数の行列の総積を計算する問題になるので有理式同士の演算が不要で実装が軽いです。)
頂点 \(1\) へ向かう辺が存在する場合
頂点 \(1\) へ向かう辺がある場合を考えます。
頂点 \(1\) への辺が存在する時、walk の構造は
\[1 \to \dots \to 1 \to \dots \to 1 \to \dots \to N\]
というように「 \(1\) に戻ってくるのを何度か繰り返した後に \(N\) に行く」という構造になっています。
この構造を利用して母関数で記述します。「頂点 \(1\) 始点, 頂点 \(1\) 終点の \(m\) 辺 walk のうち, 途中で \(1\) を通らないものの個数」の母関数を考えると、終点の直前に止まった頂点で場合分けを行うことで
\[\left(\sum_{ \lbrace n \vert D_n = 1 \rbrace } F_n\right) \times x = x \sum_{1 \leq n \leq N} D_n F_n\]
であることが確認できます。
\[G_n = \sum_{1 \leq m \leq n} {D_m F_m}\]
と置くと前述の母関数は \(x G_N\) と表せます。
次に「頂点 \(1\) 始点, 頂点 \(1\) 終点の \(m\) 辺 walk の個数」の母関数を考えると、これは 「『頂点 \(1\) 始点, 頂点 \(1\) 終点の walk のうち, 途中で \(1\) を通らないもの』を \(0\) 回以上繰り返して長さが \(m\) になるものの個数」と言えるので、
\[1 + xG_N + x^2 G_N^2 + \dots = \frac{1}{1 - x G_N}\]
となります。よってこれに 「頂点 \(1\) 始点, 頂点 \(N\) 終点の途中で \(1\) を通らない \(m\) 辺 walk の個数」の母関数 \(F_N\) を掛けた
\[\frac{F_N(x)}{1 - x G_N(x)}\]
が求めたい答えの母関数になります。
また、行列積で持つ情報を増やすことで \(G_N(x)\) を \(F_N(x)\) の計算中に同時に計算できます。(次式を参照してください)
\[ \left( \begin{array}{c} F_{n+1} \\ F_{n} \\ G_{n+1} \end{array} \right) = \left( \begin{array}{ccc} P_n Q_n & P_n R_n & 0\\ 1 & 0 & 0 \\ D_n P_n Q_n & D_n P_n R_n & 1\\ \end{array} \right) \left( \begin{array}{c} F_n \\ F_{n-1} \\ G_n \end{array} \right) \]
以上より \(F_N(x), G_N(x)\) を \(\mathrm{O}(N \log^2 N)\) で計算することができました。また、\(\lbrack x^K \rbrack \frac{F_N(x)}{1 - x G_N(x)}\) の計算は分割統治で \(\mathrm{O}(K \log^2 K)\) 、あるいは ABC260Ex 解説末尾で説明した方法 で \(\mathrm{O}(K \log K)\) で計算できます。
よってこの問題を
\(\mathrm{O}(N \log^2 N + K \log K)\) で計算することができて、これは十分高速です。(\(N\) に依存する項は定数倍に \(27\) が付く一方 \(K\) の方にはつかないので\(\mathrm{O}(N \log^2 N + K \log^2 K)\) でも高速に動作すると考えられます)
posted:
last update: