この問題は、線形漸化式によって定義される数列の第 \(N\) 項を高速に計算するアルゴリズムをテーマとする問題です。いつものように関連する話題をまとめてから解説に移ります。

繰り上がりを持つ桁 DP

はじめに、線形漸化式の 第 \(N\) 項の計算と関連深い DP を紹介します。
次の例題を考えます。

\(a_1\) 円, \(a_2\) 円, \(\dots\) \(,a_K\) 円硬貨があります。\(N\) 円支払う方法は何通りありますか？\((\text{mod }998244353)\)
\(K \leq 100, S = \sum_i a_i \leq 1000, N \leq 10^{18}\)

この例題は繰り上がりを持つ桁 DP で高速に解くことができます。簡潔に説明します。(橙 diff 程度の内容なので容易ではないかもしれません, 下の疑似コードも参考にしてください)

\(a_i\) 円硬貨を 使った枚数を \(c_i\) とします。\(c_1, \dots, c_K\) を下の bit から決めていくことを考えます。 1 つの bit に注目したときは, ある硬貨を使うか使わないかを考えればよいので \(dp[x] \gets dp[x] + dp[x-a]\) という要領の部分和 DP で計算できます。
また、\(c_1, \dots, c_K\) の \(x-1\) bit 目までを決めた時点で, 決まっていないのは \(x\) bit 目以上なので金額の総和 \(\text{mod }2^x\) は一定になり、\(x-1\) bit 目が一致していない状態は取り除いてもよいことがわかります。
結局, \(x-1\) bit 目に対応する遷移を計算し終えた時点で \(x-1\) bit 目が一致していない部分を dp テーブルから取り除いてテーブルのサイズを半分にして、残りを繰り上がりとして次の bit での計算に持ち越すという DP が動作します。繰り上がるごとに値は半分になるので DP テーブルのサイズは \(S + S/2 + S/4 + \dots \lt 2S\) になるため、\(\mathrm{O}(K S \log N)\) 程度の計算量で計算できます。

carry <- (zero-padding array of length 2S + 1)
carry[0] <- 1
while M != 0 do
  for i = 1 to K do
    for k = 2S - a_i to 0 do
      carry[k + a_i] <- carry[k + a_i] + carry[k]
    end for
  end for
  for k = 0 to 2S + 1:
    p = 2k + (N mod 2)
    if p < 2S + 1 then
      carry[k] <- carry[p]
    else 
      carry[k] <- 0
    end if
  end for
  M <- floor(M / 2)
return carry[0]

以下に類題をいくつかまとめたので理解の確認にお使いください。

類題(ARC, CF Div.1 ネタバレ注意)

• yukicoder No.137 (ほぼ同じ問題です)
• ABC271 Ex の別解
• ARC156 D
• CF Div.1 #616 F

桁 DP の母関数を用いた解釈

この章はRyuhei Mori 氏による記事を大きく参考にしました。

例題を解くのに用いた桁 DP は、形式的冪級数に対する操作に言い換えることができます。以下に詳しく説明します。

まず、この例題の数え上げる対象を母関数で表してみましょう。
\(a_i\) 円硬貨を \(0\) 枚以上使用するのに対応する母関数は

\[ 1 + x^{a_1} + x^{2 a_1} + \dots = \frac{1}{1 - x^{a_1}} \]

という形式的冪級数で記述できます。よって

\[ \lbrack x^N \rbrack \frac{1}{1 - x^{a_1}} \cdot \frac{1}{1 - x^{a_2}} \cdot \dots \cdot \frac{1}{1 - x^{a_K}} \]

が求める答えになります。ここで、

\[ \begin{aligned} \frac{1}{1-x} &= \frac{1 + x}{1 - x^2} \\ &= \frac{(1 + x)(1 + x^2)}{1 - x^4} \\ &= \frac{(1 + x)(1 + x^2)(1 + x^4)}{1 - x^8} \\ &\vdots \\ &= \prod_{b \geq 0} (1 + x^{2^b}) \end{aligned} \]

が成り立つのを利用すると、

\[ F(x) = (1 + x^{a_1})(1 + x^{a_2}) \dots (1 + x^{a_K}) \]

としたとき、求めたい答えは

\[ \lbrack x^N \rbrack \prod_{k \geq 0} F(x^{2^k}) \]

になり、\(F(x^{2^k})\) の積の形で表せることが分かりました。

このように、桁 DP で解くことができる数え上げの対象は、母関数で記述すると基数を \(q\) として (\(x^{q^0}\) の関数) \(\times\) (\(x^{q^1}\) の関数) \(\times\) (\(x^{q^2}\) の関数) \(\times \dots\) という形の式で表せることがあります。この形の式に対して \(\lbrack x^N \rbrack\) を高速に求める方法を考えてみましょう。(以下では \(q=2\) の場合を説明します。)

まず、形式的冪級数 \(B(x)\) が形式的冪級数 \(Q_0(x), Q_1(x), \dots\) を用いて

\[B(x) = \prod_{k \geq 0} Q_{k}(x^{2^k})\]

と表せる時、\(B(x)\) を 2-デジタル冪級数 と呼ぶことにします。

また、母関数の奇数次 / 偶数次 だけを取り出して新たな母関数を作る操作を考えます。\(\mathcal{S}_i\) (\(i \in \lbrace 0, 1 \rbrace\)) を次のように定義します。

\[\mathcal{S}_i \left(\sum_{n \geq 0} a_n x^n \right) = \sum_{m \geq 0} a_{2m + i} x^m\]

\(\mathcal{S}_i\) を 2-デジタル冪級数に作用させることを考えます。すると次の事実が成り立ちます。

2-デジタル冪級数 \(B_0(x), B_1(x), \dots\) が形式的冪級数 \(Q_0(x), Q_1(x), \dots\) を用いて

\[B_n(x) = \prod_{k \geq 0} Q_{n + k}(x^{2^k})\]

と表せるとする。また、\(A(x)\) を形式的冪級数とする。このとき非負整数 \(N, n\) に対して次の式が成り立つ。

\[\lbrack x^N \rbrack A(x) B_n(x) = \lbrack x^{\lfloor N/2 \rfloor} \rbrack \mathcal{S}_{N \bmod 2}(A(x) Q_n(x)) B_{n+1}(x) \]

(証明) 形式的冪級数 \(A(x)\) と \(B_n(x)\) の積に \(\mathcal{S}_i\) を作用させると、

\[ \begin{aligned} \mathcal{S}_i \left( A(x) B_n(x) \right) &= \mathcal{S}_i(A(x) Q_n(x)) \prod_{k \geq 1} Q_{n+k}(x^{2^{k-1}}) \\ &= \mathcal{S}_i(A(x) Q_n(x)) B_{n+1}(x) \end{aligned} \]

になります。また、形式的冪級数 \(C(x)\) について

\[ \lbrack x^N \rbrack C(x) = \lbrack x^{\lfloor N/2 \rfloor} \rbrack \mathcal{S}_{N \bmod 2}(C(x)) \]

が成り立つので、上式の \(C(x)\) に \(A(x) B_n(x)\) を代入して

\[ \begin{aligned} \lbrack x^N \rbrack A(x) B_n(x) &= \lbrack x^{\lfloor N/2 \rfloor} \rbrack \mathcal{S}_{N \bmod 2}(A(x) B_n(x)) \\ &= \lbrack x^{\lfloor N/2 \rfloor} \rbrack \mathcal{S}_{N \bmod 2}(A(x) Q_n(x)) B_{n+1}(x) \end{aligned} \]

を得られました。(証明終わり)

以上で証明した式は、左辺が \([x^N]\) を計算する式で、右辺が \([x^{\lfloor N / 2 \rfloor}]\) を計算する式になっているのを利用すると、以下の疑似コードのような再帰的な手順によって \(\lbrack x^N \rbrack B_0(x)\) を計算できると分かります。(式中の S(i, F) は \(\mathcal{S}_i(F)\) の意味)

A <- 1
k <- 0
while N > 0 do
  A <- S(N mod 2, A Q_k)
  k <- k + 1
  N <- floor(N / 2)
end while
return [x^0] A

計算量は \(A\) と \(Q_k(x)\) を畳み込む部分が支配的です。よって、\(B_n(x)\) を 2-デジタルな次数の低い式の積で表現できれば、このアルゴリズムは高速に動作します。

例題で求めたいのは

\[ \lbrack x^N \rbrack \prod_{k \geq 0} F(x^{2^k}) \]

でした。よって、上述のアルゴリズムに \(Q_n \gets F(x)\) を代入すれば母関数を利用した解法を導出できます。\(A\) の次数は常に \(S\) 以下なので畳み込み 1 回あたりの計算量は \(\mathrm{O}(S \log S)\) なので、アルゴリズム全体では \(\mathrm{O}(S \log S \log N)\) の計算量になります。
また、この解法を凝視すると、計算している内容が桁 DP の場合と完全に一致していることがわかります。(上から 4 行目の部分は部分和 DP + DP テーブルの半減、と同じ処理をしています。)

線形漸化式への応用

このように 桁 DP を母関数に一般化したアルゴリズムは、一般項が線形漸化式で表せる数列の第 \(N\) 項の計算に応用できます。

数列 \(a_0, a_1, a_2, \dots\) の一般項 \(a_n\) がある定数 \(c_1, \dots, c_d\) を用いて

\[a_n = \sum_{i=1}^d c_i a_{n-i}　(n \geq d)\]

と線形漸化式の形で表せる時、\((a_n)\) を 線形回帰数列 と呼びます。
線形回帰数列の第 \(N\) 項は行列累乗で \(\mathrm{O}(d^3 \log N)\) 程度で計算できることが知られていますが、母関数を利用した時間計算量 \(\mathrm{O}(d \log d \log N)\) のアルゴリズムを説明します。

線形回帰数列の母関数は有理式(多項式を分子・分母に持つ分数)で表せることが確認できます。\(a_n\) の母関数を \(A(x) = \sum_{n \geq 0} a_n x^n\) とすると、\(n \geq d\) の範囲で

\[ \begin{aligned} \lbrack x^n \rbrack \left(1 - \sum_{i = 1}^d c_i x^i\right)A(x) &= a_n - \sum_{i=1}^d c_i a_{n - i} \\ &= 0 \end{aligned} \]

が成り立つので、

\[Q(x) = 1 - \sum_{i = 1}^d c_i x^i\]

とすると \(A(x)Q(x)\) は \(d-1\) 次以下の多項式になることがわかります。

\[ P(x) = \left(\sum_{n=0}^{d-1} a_n x^n \right) Q(x) \bmod{x^d} \]

とおくと \(P(x) \equiv A(x) Q(x) \pmod{x^d}\) が成り立つので、\(A(x)\) と \(P(x), Q(x)\) の間には

\[A(x) = \frac{P(x)}{Q(x)}\]

という関係式が成り立つのが分かります。

よって、有理式 \(\frac{P(x)}{Q(x)}\) \((\deg(P) \lt \deg(Q))\) の \(N\) 次の係数が計算できれば線形回帰数列の第 \(N\) 項が求まります。この計算は Bostan-Mori 法 と呼ばれるアルゴリズムが知られているので紹介します。

線形回帰数列 \(a_0, a_1, a_2, \dots\) の要素を偶数番目と奇数番目で分けると、

• 偶数番目だけ取り出した数列：\(a_0, a_2, a_4, \dots\)
• 奇数番目だけ取り出した数列：\(a_1, a_3, a_5, \dots\)

の 2 種類の数列ができます。実は、これらの数列はどちらも高々 \(d\) 次の線形回帰数列になります(後に確認できます)。
また、\((a_n)\) の母関数を \(A(x)\) としたとき、この 2 つの数列の母関数は、\(\mathcal{S}\) を使うとそれぞれ \(\mathcal{S}_0(A(x)), \mathcal{S}_1(A(x))\) と表せます。

よって、\(A(x)\) から \(\mathcal{S}_i(A(x))\) を高速に計算できれば、

\[\lbrack x^N \rbrack A(x) = \lbrack x^{\lfloor N/2 \rfloor} \rbrack \mathcal{S}_{N \bmod 2} (A(x))\]

という関係式を利用して求めたい項を半減させていく再帰で答えを計算できます。

\(\mathcal{S}_i\left(\frac{P(x)}{Q(x)}\right)\) は以下の手順で計算できます。\(\frac{P(x)}{Q(x)}\) の分子と分母に \(Q(-x)\) を掛けて

\[\frac{P(x)}{Q(x)} = \frac{P(x)Q(-x)}{Q(x)Q(-x)}\]

と変形すると、分母の \(Q(x)Q(-x)\) は偶多項式なのである多項式 \(T(x)\) を用いて

\[T(x^2) = Q(x)Q(-x)\]

と表せます。この \(T\) を使うと、

\[ \begin{aligned} \mathcal{S}_i\left( \frac{P(x)}{Q(x)} \right) &= \mathcal{S}_i\left( \frac{P(x)Q(-x)}{T(x^2)} \right) \\ &= \frac{\mathcal{S}_i (P(x) Q(-x))}{T(x)} \end{aligned} \]

になるのがわかります。よって \(\mathcal{S}_i(A(x))\) に対応する有理式は、\(P(x)\) と \(Q(-x)\), \(Q(x)\) と \(Q(-x)\) を畳み込めば \(\mathrm{O}(d \log d)\) で計算できるとわかります。

以上より、有理式 \(\frac{P(x)}{Q(x)}\) の \(N\) 次の係数は (\(\deg(Q) = d\) として) 次のアルゴリズムで \(\mathrm{O}(d \log d \log N)\) で計算できます。

input : P(x), Q(x), N
output : [x^N] P(x) / Q(x)

while N > 0 do
  P(x) <- S(N mod 2, P(x) Q(-x))
  Q(x) <- S(0, Q(x) Q(-x))
  N <- floor(N / 2)
end while
return P(0) / Q(0)

そして、このアルゴリズムは前章で説明したアルゴリズムの一種としても説明できます。すなわち、\((Q_n(x))_{n \geq 0}\) を \(Q_0(x) = Q(x), Q_{n+1}(x^2) = Q_n(x) Q_n(-x)\) を満たす多項式の列としておくと

\[ \begin{aligned} \frac{P(x)}{Q(x)} &= \frac{P(x)Q_0(-x)}{Q_0(x)Q_0(-x)} \\ &= P(x)Q_0(-x) \times \frac{1}{Q_1(x^2)} \\ &= P(x)Q_0(-x) \times \frac{Q_1(-x^2)}{Q_1(x^2)Q_1(-x^2)} \\ &= P(x)Q_0(-x) \times Q_1(-x^2) \times \frac{1}{Q_2(x^4)} \\ &\vdots \\ &= P(x) Q_0(-x) \times \prod_{n \geq 1} Q_n(-x^{2^n}) \end{aligned} \]

になるため \(\frac{P(x)}{Q(x)}\) は 2-デジタル冪級数です。そして、前章のアルゴリズムをそのまま適用すると上のアルゴリズムと一致します。

問題の解法

元の問題の解法に移ります。\(K\)-bonacci 数列は線形回帰数列なので、元の問題を母関数に言い換えると次のようになります。

• \(\deg(P) \lt \deg(Q) = K\) であるような有理式 \(\frac{P(x)}{Q(x)}\) が与えられる。\(m \text{ AND } N = m\) であるような \(m\) に対して \(\lbrack x^m \rbrack \frac{P(x)}{Q(x)}\) の和を計算せよ。

2 通りの解法を説明します。
1 つは Bostan-Mori 法を応用した解法です。
今までの話をまとめると、Bostan-Mori 法は次の操作を \(i\) 昇順に行うアルゴリズムです。(解法から見ている方へ：ここで \(\mathcal{S}_i\) は母関数の偶数次/奇数次だけを取り出して新たな母関数を作る操作とします。つまり \(\mathcal{S}_i \left(\sum_{n \geq 0} a_n x^n \right) = \sum_{m \geq 0} a_{2m + i} x^m\) です)

• \(N\) の i bit 目が \(0\) ならば \(A(x) \gets \mathcal{S}_0(A(x))\)
• \(N\) の i bit 目が \(1\) ならば \(A(x) \gets \mathcal{S}_1(A(x))\)

元の問題の「\(m \text{ AND }N = m\) を満たす \(m\) 」という部分は、「\(N\) の \(i\) bit 目が \(1\) ならば \(m\) の \(i\) bit 目は \(0\) か \(1\), \(N\) の i bit 目が \(0\) ならば \(0\) を選べる」と言い換えられます。よって、上の場合分けを

• \(N\) の i bit 目が \(0\) ならば \(A(x) \gets \mathcal{S}_0(A(x))\)
• \(N\) の i bit 目が \(1\) ならば \(A(x) \gets \mathcal{S}_0(A(x)) + \mathcal{S}_1(A(x))\)

して、あとは Bostan-Mori 法と同様の再帰を行えばよいとわかります。

もう 1 つは「桁 DP の母関数を用いた解釈」で登場したアルゴリズムを利用した解法です。求めたい答えは, \(K\)-bonacci 数列の母関数を \(F(x)\) として

\[\lbrack x^N \rbrack F(x) \left(\sum_{m \text{ AND } N = m} x^m\right)\]

と表すことができます。\(F(x)\) は有理式なので 2-デジタルで、シグマで表される項も

\[\prod_{2^b \text{ AND } N = 2^b} \left(1 + x^{2^b} \right)\]

という形で表されるので 2-デジタルです。よって「桁 DP の母関数を用いた解釈」の \(Q_n(x)\) に適切に式を代入するアルゴリズムが動作します。

実装するとわかりますが、2 つの解法は結果的に同じアルゴリズムに帰着します。計算量は \(\mathrm{O}(K \log K \log N)\) で、十分高速です。

• 実装例(C++)

#include <iostream>
#include <vector>
using namespace std;

#include "atcoder/convolution.hpp"
#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;

int main() {
  long long K, N;
  cin >> K >> N;
  vector<mint> Q(K + 1, -1);
  Q[0] = 1;
  vector<mint> P = atcoder::convolution(Q, vector<mint>(K, 1));
  P.resize(K);
  for (; N; N >>= 1) {
    vector<mint> Q2{Q};
    for (int i = 1; i <= K; i += 2) Q2[i] = -Q2[i];
    P = atcoder::convolution(P, Q2);
    Q = atcoder::convolution(Q, Q2);
    vector<mint> S(K), T(K + 1);
    for (int i = 0; i < K; i++) S[i] = P[i * 2] + (N % 2 ? P[i * 2 + 1] : 0);
    for (int i = 0; i < K + 1; i++) T[i] = Q[i * 2];
    P = S, Q = T;
  }
  cout << (P[0] * (Q[0].inv())).val() << "\n";
}