この間、フィボナッチ数を計算する記事を書いていたら、@fetburner氏にこういう問題を教えて頂いた：
T: フィボナッチ – Typical DP Contest | AtCoder
フィボナッチ数列の各項は直前の2項（隣接3項間の漸化式）によって定まる\[a_{n+2}=a_n+a_{n+1}\]が、それを「直前の\(K\)項」に一般化したものを高速に計算しなさい、ということだそうだ。\begin{align*}a_1&=a_2=\cdots=a_K=1,\\a_i&=\underbrace{a_{i-1}+\cdots+a_{i-K}}_{K\text{項}}\quad(i>K)\end{align*}
競技プログラミングにはあまり興味はない（制限時間があって他の人と競って〜という形式は好きではない）が、フィボナッチ数列の計算について熱弁した後に「わからん！w」と言って逃げるわけにはいかないので、この問題を解いてみることにした。結果、正答（AC）を得るまでに3時間ほどかかった。
この記事は普段競プロをしない筆者が書いたものなので、競プロしない人も安心してお読みください。競プロ特有の用語は登場しません。
（コード例はHaskellだが、「matMulやmatPowという関数が定義されているからこのプログラムでは行列乗算・行列のべき乗を使っているんだな」という程度の情報が読み取れれば十分だろう。）
前提となる記事：
今回のソースコードは
に置いた。（競プロの問題一つにGitリポジトリーを用意するのは大げさな気もするが、気にしない）
目次
大雑把な方針：行列乗算
まず、問題では数列のN番目\(a_N\)をピンポイントで計算することが求められている。こういう問題に適した方法は、フィボナッチ数列の計算で言えば「一般項（ビネの公式）を使う方式」「行列乗算を使う方式」「行列乗算を改良した方式（fast doublingを含む）」が該当する。
一般項だが、フィボナッチ数列のような既知の数列だと一般項はよく知られているのに対し、その一般化の場合はまず一般項を考えるところから始めないといけない。そして、その一般項がうまい具合に\(K\)で表されるとは限らない。何より、フィボナッチ数の計算では一般項を使うよりも行列演算の改良版（fast doubling）の方が早かったので、一般項を使う理由がない。
というわけで、この問題は「行列乗算およびその改良版」を使って解くことになる。
フィボナッチ数の場合の関係式を行列で書けば\[\begin{pmatrix}
F_{n}\\F_{n+1}
\end{pmatrix}=
\begin{pmatrix}
0&1\\
1&1
\end{pmatrix}
\begin{pmatrix}
F_{n-1}\\F_{n}
\end{pmatrix}\]となったが、直前\(K\)項の場合は\[\begin{pmatrix}
a_{i-K+1}\\a_{i-K+2}\\\vdots\\a_{i-2}\\a_{i-1}\\a_i
\end{pmatrix}=
\begin{pmatrix}
0&1&0&\cdots&0&0\\
0&0&1&&0&0\\
\vdots&&&\ddots&&\vdots\\
0&0&0&&1&0\\
0&0&0&\cdots&0&1\\
1&1&1&\cdots&1&1
\end{pmatrix}
\begin{pmatrix}
a_{i-K}\\a_{i-K+1}\\\vdots\\a_{i-3}\\a_{i-2}\\a_{i-1}
\end{pmatrix}\]となる。この行列\[\begin{pmatrix}
0&1&0&\cdots&0&0\\
0&0&1&&0&0\\
\vdots&&&\ddots&&\vdots\\
0&0&0&&1&0\\
0&0&0&\cdots&0&1\\
1&1&1&\cdots&1&1
\end{pmatrix}\]を\(T_K\)と置くことにしよう。\(T_K\)の\(N\)乗\(T_K^N\)を高速に計算できれば、\(a_N\)を高速に計算できることになる。具体的には
- \(1\le N\le K\)なら \(a_N=1\) であり、
- \(K<N\)なら\(T_K^{N-K}\)を計算し、\[\begin{pmatrix}
a_{N-K+1}\\a_{N-K+2}\\\vdots\\a_{N-1}\\a_N
\end{pmatrix}=T_K^{N-K}\begin{pmatrix}
1\\1\\\vdots\\1\\1
\end{pmatrix}\]によって \(a_N\) を求める
となる。
Haskellによる実装
筆者にはHaskellが慣れているのでHaskellを使う。日常的に競技プログラミングをするならもっと高速な言語を使えるようにするべきだろうが、単発で問題を解くだけなら使い慣れた言語を使うのが良いだろう。
まず、仕様通りに入力を読み取るコードを書く。
main :: IO ()
main = do
l <- getLine
let [(k, l')] = (reads :: ReadS Int) l
[(n, _)] = (reads :: ReadS Int) l'
putStrLn $ "K = " ++ show k ++ ", N = " ++ show n
実行してみる。
2 10 K = 2, N = 10
無事に K と N を解釈できているようだ。あとはロジックを実装すれば良い。
漸化式を順番に計算する
フィボナッチ数列の愚直なHaskell実装は
fib = 1 : 1 : zipWith (+) fib (tail fib)
だった。これを直前k項に一般化すると、
seq = replicate k 1 ++ map (sum . take k) (tails seq)
となる。この問題では mod 1000000007 で考えることを踏まえると、完全なソースコードは
import Data.Int (Int64)
import Data.List (foldl', tails)
modulo = 1000000007 :: Int64
addMod x y = (x + y) `mod` modulo
sumMod = foldl' addMod 0
main :: IO ()
main = do
l <- getLine
let [(k, l')] = (reads :: ReadS Int) l
[(n, _)] = (reads :: ReadS Int) l'
let seq = replicate k 1 ++ map (sumMod . take k) (tails seq)
print $ seq !! (n - 1)
となる。試しに K=2, N=10 を与えると 55 が、K=3, N=10 を与えると 105 が出力される。
一方、問題の条件としては \(K\le 1000\), \(N\le 10^9\), 時間制限は 2秒 となっている。大きな K, N でどのくらいの時間がかかるか確かめておこう。
- K=100, N=10000: 0.04秒（\(a_N=473837614\) ただし mod 1000000007 で。以下同様）
- K=100, N=100000: 0.21秒（\(a_N=195246043\)）
- K=100, N=1000000: 2.36秒（\(a_N=507299536\)）
- K=1000, N=10000: 0.17秒（\(a_N=636507430\)）
- K=1000, N=100000: 1.79秒（\(a_N=519884139\)）
- K=1000, N=1000000: 18.76秒（\(a_N=22223631\)）
\(N=10^9\) どころか、 \(N=10^5\) を超えるともうダメ（2秒を超える）のようだ。行列演算など、 log N で済む方法を試そう。
（漸化式で愚直に計算する方法がこの問題で実用的ではないことは分かり切っていたが、他の方法を実装するにあたって正しく実装できているかを検証する目的では役に立つ。）
行列乗算による実装
まずは、素朴に行列乗算を使って実装してみよう。行列は Data.Array を使って、 Array (Int,Int) a と表現することにする。
数列の値は mod 1000000007 で計算すること、となっているので、 mod 1000000007 で計算してくれるデータ型 N を定義した（名前が雑である）。
{-# LANGUAGE BangPatterns #-}
import Data.Array
import Data.Ix (range)
import Data.Int (Int64)
---
modulo = 1000000007 :: Int64
newtype N = N { getModuloN :: Int64 } deriving (Eq)
instance Num N where
N x + N y = N ((x + y) `mod` modulo)
N x - N y = N ((x - y) `mod` modulo)
N x * N y = N (x * y `mod` modulo)
fromInteger n = N (fromInteger (n `mod` fromIntegral modulo))
abs = undefined
signum = undefined
instance Show N where
show (N x) = show x
---
matMul :: (Num a) => Array (Int,Int) a -> Array (Int,Int) a -> Array (Int,Int) a
matMul a b = let ((i0,j0),(ix,jx)) = bounds a
((j'0,k0),(j'x,kx)) = bounds b
in if jx - j0 == j'x - j'0
then array ((i0,k0),(ix,kx))
[ ((i,k), sum [a!(i,j) * b!(j',k) | (j,j') <- zip (range (j0,jx)) (range (j'0,j'x))])
| i <- range (i0,ix)
, k <- range (k0,kx)
]
else error "Matrix size mismatch"
matPow :: (Num a) => Int -> Array (Int,Int) a -> Int -> Array (Int,Int) a
matPow k m 0 = array ((1,1),(k,k)) $
[((i,j), if i == j then 1 else 0) | i <- [1..k], j <- [1..k]]
matPow _ m i = loop (i-1) m m
where
loop 0 !_ acc = acc
loop 1 m acc = m `matMul` acc
loop i m acc = case i `quotRem` 2 of
(j,0) -> loop j (m `matMul` m) acc
(j,_) -> loop j (m `matMul` m) (acc `matMul` m)
main :: IO ()
main = do
l <- getLine
let [(k, l')] = (reads :: ReadS Int) l
[(n, _)] = (reads :: ReadS Int) l'
if n <= k
then print 1
else do
let t = array ((1,1),(k,k)) $
[((i,j), if i + 1 == j then 1 else 0) | i <- [1..k-1], j <- [1..k]]
++ [((k,j),1) | j <- [1..k]]
let u = matPow k t (n - k) :: Array (Int,Int) N
print $ getModuloN $ sum [u!(k,j) | j <- [1..k]]
実行してみると
- K=100, N=10000: 10.3秒（比較：漸化式を普通に計算する方法では0.04秒）
って……めちゃくちゃ遅いんだが〜〜〜？？？
愚直に「漸化式を順番に計算する」方法よりも遅くなっている。これはHaskellの遅延評価の罠を踏んでいる可能性が大である。試しに、unboxed array (Data.Array.Unboxed) を使ってみよう。
ここで悲報だが、unboxed arrayはnewtypeと相性が悪い。Haskellではnewtypeで新しい型を定義することによって Ord や Num 等のインスタンスを上書きする、ということがよく行われていて、実際上記のコードでは newtype N = N Int64 と新しい型を定義し、それによって「行列演算」と「mod 1000000007 での計算」を分離できている（matMul や matPow は汎用的に使える）。
しかし、unboxed arrayに関しては、 Int64 を要素にもつunboxed arrayは作れるが、そのnewtypeである newtype N = N Int64 を要素にもつunboxed arrayは作れない（UArray i Int64 はOKだが、 UArray i Nはダメ）。
というわけで、行列の乗算とべき乗の関数をそれぞれ「Int64 のunboxed arrayを扱い」「要素を mod 1000000007 で計算する」ように書き換える。この際、間違えて要素をmodで計算せずに Int64 の + や * を使ってしまってもエラーは起こらない。抽象化を愛するHaskellユーザーとしてはこのような汎用性のないコードを書くことになって非常に悔しいが……。
{-# LANGUAGE BangPatterns #-}
module MatMulUnboxed where
import Data.Array.Unboxed
import Data.Ix (range)
import Data.Int (Int64)
import Data.List (foldl')
---
modulo = 1000000007 :: Int64
addMod x y = (x + y) `mod` modulo
mulMod x y = (x * y) `mod` modulo
sumMod = foldl' addMod 0
---
matMul :: UArray (Int,Int) Int64 -> UArray (Int,Int) Int64 -> UArray (Int,Int) Int64
matMul a b = let ((i0,j0),(ix,jx)) = bounds a
((j'0,k0),(j'x,kx)) = bounds b
in if jx - j0 == j'x - j'0
then array ((i0,k0),(ix,kx))
[ ((i,k), sumMod [(a!(i,j)) `mulMod` (b!(j',k)) | (j,j') <- zip (range (j0,jx)) (range (j'0,j'x))])
| i <- range (i0,ix)
, k <- range (k0,kx)
]
else error "Matrix size mismatch"
matPow :: Int -> UArray (Int,Int) Int64 -> Int -> UArray (Int,Int) Int64
matPow k m 0 = array ((1,1),(k,k)) $
[((i,j), if i == j then 1 else 0) | i <- [1..k], j <- [1..k]]
matPow _ m i = loop (i-1) m m
where
loop 0 !_ acc = acc
loop 1 m acc = m `matMul` acc
loop i m acc = case i `quotRem` 2 of
(j,0) -> loop j (m `matMul` m) acc
(j,_) -> loop j (m `matMul` m) (acc `matMul` m)
main :: IO ()
main = do
l <- getLine
let [(k, l')] = (reads :: ReadS Int) l
[(n, _)] = (reads :: ReadS Int) l'
if n <= k
then print 1
else do
let t = array ((1,1),(k,k)) $
[((i,j), if i + 1 == j then 1 else 0) | i <- [1..k-1], j <- [1..k]]
++ [((k,j),1) | j <- [1..k]]
let u = matPow k t (n - k) :: UArray (Int,Int) Int64
print $ sumMod [u!(k,j) | j <- [1..k]]
やってみよう。
- K=100, N=10000: 0.70秒（\(a_N=473837614\)）
- K=100, N=100000: 0.79秒（\(a_N=195246043\)）
- K=100, N=1000000: 0.99秒（\(a_N=507299536\)）
- K=100, N=10000000: 1.10秒（\(a_N=773373988\)）
- K=100, N=100000000: 1.34秒（\(a_N=716893641\)）
- K=100, N=1000000000: 1.52秒（\(a_N=103057669\)）
- K=1000, N=10000: ………
K=100 の場合は、さすが \(O(\log N)\) のアルゴリズムと言うべきか、 N が \(10^9\) 程度まで大きくなっても制限時間内に計算ができている。
しかし、 K=1000 となると途端に計算ができなくなる。これは、素朴な行列乗算のアルゴリズムが \(O(K^3)\) で、K=100からK=1000になると計算量が1000倍になることが関係していると考えられる。
とにかく、普通に行列を計算していたのではダメだ。
さらなる高速化
この間の記事「最速のフィボナッチ数計算を考える」では、\[\begin{pmatrix}0&1\\1&1\end{pmatrix}^n=\begin{pmatrix}F_{n-1}&F_n\\F_n&F_{n-1}+F_n\end{pmatrix}=\begin{pmatrix}F_{n+1}-F_n&F_n\\F_n&F_{n+1}\end{pmatrix}\]という関係式を使って、2×2行列の代わりに数の2つ組（その記事では FibPair と呼んだ）とその演算を考えた。一般の \(K\) の場合にも同様の関係式があれば、普通に行列で計算するよりも高速に計算ができそうである。
K=3 の場合の観察
フィボナッチ数の場合は、実際にいくつか行列の\(n\)乗を計算することで法則性が見えてきた。\(K=3\)の場合、つまり\(T_3=\begin{pmatrix}0&1&0\\0&0&1\\1&1&1\end{pmatrix}\)の\(n\)乗を実際にいくつか計算してみよう。\begin{align*}
T_3^0&=\begin{pmatrix}1&0&0\\0&1&0\\0&0&1\end{pmatrix},\\
T_3^1&=\begin{pmatrix}0&1&0\\0&0&1\\1&1&1\end{pmatrix},\\
T_3^2&=\begin{pmatrix}0&0&1\\1&1&1\\1&2&2\end{pmatrix},\\
T_3^3&=\begin{pmatrix}1&1&1\\1&2&2\\2&3&4\end{pmatrix},\\
T_3^4&=\begin{pmatrix}1&2&2\\2&3&4\\4&6&7\end{pmatrix},\\
T_3^5&=\begin{pmatrix}2&3&4\\4&6&7\\7&11&13\end{pmatrix},\\
T_3^6&=\begin{pmatrix}4&6&7\\7&11&13\\13&20&24\end{pmatrix}
\end{align*}勘の良い読者なら法則性が見えてきただろう。どうやら、\(\alpha_n\), \(\beta_n\), \(\gamma_n\) を適当に置けば、\(T_3^n\)は\[T_3^n=\begin{pmatrix}\alpha_n&\beta_n&\gamma_n\\\gamma_n&\alpha_n+\gamma_n&\beta_n+\gamma_n\\\beta_n+\gamma_n&\beta_n+2\gamma_n&\alpha_n+\beta_n+2\gamma_n\end{pmatrix}\]と表されるようだ。書き方を変えれば、\[T_3^n=\alpha_n\begin{pmatrix}1&0&0\\0&1&0\\0&0&1\end{pmatrix}+\beta_n\begin{pmatrix}0&1&0\\0&0&1\\1&1&1\end{pmatrix}+\gamma_n\begin{pmatrix}0&0&1\\1&1&1\\1&2&2\end{pmatrix}\]となる。この右辺に出てきた行列は\(I\)（単位行列）, \(T_3\), \(T_3^2\) に他ならない：\[T_3^n=\alpha_n\cdot I+\beta_n\cdot T_3+\gamma_n\cdot T_3^2\]
フィボナッチ数列（直前2項の和で表される数列）の計算において \(\begin{pmatrix}0&1\\1&1\end{pmatrix}^n\) が実際には4つも成分を必要とせず、2つの成分で表される（2次元）のと同じように、直前3項の和で表される数列の計算においては \(T_3^n\) は3つの成分で表される（3次元）のだ。
行列と多項式
これで\(K=2\)と\(K=3\)の場合に\begin{align*}T_2^n&=F_{n-1}\cdot I+F_n\cdot T_2,\\T_3^n&=\alpha_n\cdot I+\beta_n\cdot T_3+\gamma_n\cdot T_3^2\end{align*}という風に、\[T_K^n\in\langle I, T_K, T_K^2, \ldots, T_K^{K-1}\text{によって張られる線型空間}\rangle\]となることがわかった。この関係は一般に成り立つだろうか？
成り立つ。ケーリー・ハミルトンの定理によると、\(K\)次の正方行列\(T\)に対して\(K\)次の多項式\(P(X)=X^K+a_{K-1}X^{K-1}+\cdots+a_1X+a_0\)が存在し、\(P(T)=0\)となる。
ということは関係式\(T^K+a_{K-1}\cdot T^{K-1}+\cdots+a_1\cdot T+a_0\cdot I=0\)が成り立ち、大きな\(n\)に対しての\(T^n\)はこの関係式を繰り返し使っていけば次数を下げて、\(I, T,\ldots,T^{K-1}\)の線型結合で表すことができる。
この「関係式を使って次数を下げる」という操作には「多項式除算」という名前がついていて、その言葉を使えば「多項式 \(X^n\) を \(P(X)\) で割った余りを求め、それに \(T\) を代入する」となる。
（宣伝だが、筆者が以前書いた「週刊 代数的実数を作る」の「#6 代数的数の演算 まとめ」では、代数的数の \(n\) 乗を計算する際に「\(X^n\) を最小多項式 \(P(X)\) で割る」ということをやっている。行列と代数的数、対象としては違うように見えるが、どちらも最小多項式を持つなど、似た側面がある。）
さて、行列 \(T\) に対して \(P(T)=0\) となるような多項式\(P\)はどうやって求めたら良いか。一般の行列に対してこのような多項式を求めるのはちょっと大変だが、今回の行列 \(T=T_K\) は線型漸化式に付随する行列なので、かなりよくわかっている。
実際、行列\(T_K\)は\[a_{n+K}=a_{n+K-1}+\cdots+a_{n}\]という線形漸化式に付随する行列なので、\[T_K^K=T_K^{K-1}+\cdots+T_K+I\]という関係式を満たす（\(P(T)=0\) となる多項式として \(P(X)=X^K-X^{K-1}-\cdots-X-1\) を取れる）。
（例：\(T_2=\begin{pmatrix}0&1\\1&1\end{pmatrix}\) は \(T_2^2=T_2+I\) を満たし、\(T_3=\begin{pmatrix}0&1&0\\0&0&1\\1&1&1\end{pmatrix}\) は \(T_3^3=T_3^2+T_3+I\) を満たす。）
結局、アルゴリズムとしては
- \(1\le K\le N\) なら \(a_N=1\) である（終了）。
- \(K<N\) なら、
- 多項式 \(X^{N-K}\) を多項式 \(P(X)=X^K-X^{K-1}-\cdots-X-1\) で割った余りを計算し、 \(F(X)\) とおく。
- \(X^{N-K} \bmod P(X)\) の計算には、多項式に関するpowModみたいなことをする。
- \(F\) に対して行列 \(T_K\) を代入したもの \(F(T_K)\) を計算し、\(\begin{pmatrix}a_{N-K+1}\\a_{N-K+2}\\\vdots\\a_N\end{pmatrix}=F(T_K)\begin{pmatrix}
a_1\\a_2\\\vdots\\a_K
\end{pmatrix}=F(T_K)\begin{pmatrix}
1\\1\\\vdots\\1
\end{pmatrix}\) により \(a_N\) を計算する。
- 多項式 \(X^{N-K}\) を多項式 \(P(X)=X^K-X^{K-1}-\cdots-X-1\) で割った余りを計算し、 \(F(X)\) とおく。
となる。
なお、ここまでくれば行列を使う必要はもはやなく、\(T_K\) を「漸化式 \(a_{n+K}=a_{n+K-1}+\cdots+a_n\) を満たす数列 \((a_n)\) をシフトさせる関数 \((a_n)\mapsto (a_{n+1})\)」として実装して良い。
また、多項式除算は割る側の多項式の形が \(P(X)=X^K-X^{K-1}-\cdots-X-1\) だと分かっているので、割とハードコーディングできる。
Haskell実装
多項式の表現には Data.Vector の Vector 型を使う（「週刊 代数的実数を作る」でも多項式を Vector で表した）。係数の列を昇べきの順で保持する。多項式演算の全てを実装するのは面倒なので、掛け算、乗算、\(P(X)\) による余りの計算を実装する。
実装例：
{-# LANGUAGE BangPatterns #-}
import Data.Int (Int64)
import qualified Data.Vector as V
import Data.List (tails)
modulo = 1000000007 :: Int64
newtype N = N { getModuloN :: Int64 } deriving (Eq)
instance Num N where
N x + N y = N ((x + y) `mod` modulo)
N x - N y = N ((x - y) `mod` modulo)
N x * N y = N (x * y `mod` modulo)
fromInteger n = N (fromInteger (n `mod` fromIntegral modulo))
abs = undefined
signum = undefined
instance Show N where
show (N x) = show x
-- 多項式は
-- V.fromList [a,b,c,...,z] = a + b * X + c * X^2 + ... + z * X^(k-1)
-- により表す。
-- 多項式を X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
reduce :: (Eq a, Num a) => Int -> V.Vector a -> V.Vector a
reduce k v | V.last v == 0 = V.init v
| V.length v <= k = v
| otherwise = let b = V.last v
l = V.length v
in reduce k (V.imap (\i a -> if i >= l - k - 1 then a + b else a) (V.init v))
-- 多項式の積を X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
mulP :: (Eq a, Num a) => Int -> V.Vector a -> V.Vector a -> V.Vector a
mulP k !v !w = reduce k $ V.generate (V.length v + V.length w - 1) $
\i -> sum [(v V.! (i-j)) * (w V.! j) | j <- [0..V.length w-1], j <= i, j > i - V.length v]
-- 不定元
ind :: (Num a) => V.Vector a
ind = V.fromList [0,1]
-- 多項式の（mod X^k - X^(k-1) - ... - X - 1 での）べき乗
powP :: (Eq a, Num a) => Int -> V.Vector a -> Int -> V.Vector a
powP k _ 0 = V.fromList [1]
powP k m i = loop (i-1) m m
where
loop 0 !_ !acc = acc
loop 1 m acc = mulP k m acc
loop i m acc = case i `quotRem` 2 of
(j,0) -> loop j (mulP k m m) acc
(j,_) -> loop j (mulP k m m) (mulP k acc m)
main :: IO ()
main = do
l <- getLine
let [(k, l')] = (reads :: ReadS Int) l
[(n, _)] = (reads :: ReadS Int) l'
if n <= k
then print 1
else do
let f = powP k ind (n - k) -- X^(n-k) mod X^k - X^(k-1) - ... - X - 1
let seq = replicate k 1 ++ map (sum . take k) (tails seq) -- 数列
print $ getModuloN $ sum $ zipWith (*) (V.toList f) (drop (k-1) seq)
最後の3行について説明しておこう。まず、f は \(X^{N-K}\bmod X^K-\cdots-X-1\) の係数を昇べきの順に並べたもの、つまり\begin{align*}F(X)&=X^{N-K}\bmod X^K-\cdots-X-1\\&=f_0+f_1X+\cdots+f_{K-1}X^{K-1}\end{align*}とおいたときに \(\mathtt{f}=[f_0,f_1,\ldots,f_{K-1}]\) である。
次に、 seq は求めたい数列 \(a_n\) を愚直に「漸化式を順番に計算」することによって計算する無限リストである。例えば \(K=2\) の場合は seq=[1,1,2,3,5,8,...] という無限リストとなるし、 \(K=3\) の場合は seq=[1,1,1,3,5,9,17,31,...] という無限リストとなる。
そして、先ほど述べたアルゴリズムの式を、 \(f_i\) などを使って表すと\begin{align*}
\begin{pmatrix}a_{N-K+1}\\a_{N-K+2}\\\vdots\\a_N\end{pmatrix}
&=F(T_K)\begin{pmatrix}
a_1\\a_2\\\vdots\\a_K
\end{pmatrix}\\
&=(f_0\cdot I+f_1\cdot T_K+\cdots+f_{K-1}\cdot T_K^{K-1})\begin{pmatrix}
a_1\\a_2\\\vdots\\a_K
\end{pmatrix}\\
&=f_0 \begin{pmatrix}
a_1\\a_2\\\vdots\\a_K
\end{pmatrix}+f_1 \begin{pmatrix}
a_2\\a_3\\\vdots\\a_{K+1}
\end{pmatrix}+\cdots+f_{K-1} \begin{pmatrix}
a_K\\a_2\\\vdots\\a_{2K-1}
\end{pmatrix}
\end{align*}となる。特に、\[a_N=f_0\cdot a_K+f_1\cdot a_{K+1}+\cdots+f_{K-1}\cdot a_{2K-1}\]なので、\(f_i\) と「数列 \(a_n\) のはじめの \(K-1\) 項を取り除いたもの」の内積みたいなものを sum . zipWith (*) で計算してやれば \(a_N\) が求まる。
このプログラムで計算してみると、
- K=100, N=10000: 0.02秒（\(a_N=473837614\)）
- K=100, N=100000: 0.04秒（\(a_N=195246043\)）
- K=100, N=1000000: 0.05秒（\(a_N=507299536\)）
- K=100, N=10000000: 0.05秒（\(a_N=773373988\)）
- K=100, N=100000000: 0.08秒（\(a_N=716893641\)）
- K=100, N=1000000000: 0.09秒（\(a_N=103057669\)）
- K=1000, N=10000: 1.76秒（\(a_N=636507430\)）
- K=1000, N=100000: 3.55秒（\(a_N=519884139\)）
- K=1000, N=1000000: 7.62秒（\(a_N=22223631\)）
- K=1000, N=10000000: 8.56秒（\(a_N=156587014\)）
- K=1000, N=100000000: 14.19秒（\(a_N=996909980\)）
- K=1000, N=1000000000: 22.34秒（\(a_N=289587781\)）
となった（並列化された上でのreal timeなので、同じマシンでもシングルスレッドだともうちょっと時間がかかるかもしれない）。
パラメーターの上限である \(K=1000\), \(N=10^9\) でも一応計算できてはいるのだが、22秒と、時間制限はパスできていない。アルゴリズム自体はそこそこ良いもののはずなので、あとは……そう、Haskellの遅延評価をどうにかすれば高速化して時間制限をパスできる可能性がある。
さっきはunboxed arrayを使ったが、今度はVectorを使っているので使うのはunboxed vectorである。そして、Data.Vector.Unboxedにあるunboxed vectorもnewtypeとの相性が悪い！Haskellはクソ！遅延評価は悪い文明！
コードは省略する（GitHubを参照）が、プログラムを Data.Vector.Unboxed.Vector Int64 を使うように書き換えたところ、
- K=1000, N=1000000000（10の9乗）: 1.47秒
で計算できるようになった。これならアルゴリズムにこれ以上改良を加えなくても、時間制限をパスできそうである。
実際に提出してみると、
提出 #4147694 – Typical DP Contest
というわけで、実行時間1759msで通った。
他の人の解答（C++とか）を見ると数百msで通している人が多いので、Haskellでももうちょっと高速化する余地はあるのかもしれないが、今回はここで良しとしておく。
【2月4日 追記】フィボナッチ数計算のfast doubling相当（繰り返し乗算の際に二進表記を「左から右へ」辿り、「X倍してmod P(X)」の部分を特化した関数に置き換える）の方法を実装したところ、1130msまで縮んだ。前回の記事でfast doublingと呼ばれる方法のエッセンスを抽象的な形で記述しておいたからこそ、数列が変わっても同じ手法を適用できたわけだ（さもなくば、今回の数列の \(a_n\) と \(a_{2n}\) の間に成り立つ漸化式を見つけるところから始めなくてはならなかっただろう）。抽象化万歳！
これ以上高速化したかったら、多項式乗算＆剰余計算の部分をもうちょっとどうにかした方が良いだろう。
解いてみたまとめ
Haskellの遅延評価はクソ。unboxed arrayやunboxed vectorとnewtypeの相性が悪いのもクソ。競プロには向いていない。
競プロ云々を抜きにしても、unboxed arrayやunboxed vectorとnewtypeの相性が悪いのはどうにかならんのか？と思ったのでnewtypeとの相性が良いunboxed vector（のラッパー）を作った→https://github.com/minoki/unboxing-vector そのうちHackageに上げて記事も書きたい…。
今回使ったアルゴリズムの要点は、行列のべき乗を計算する際にその行列の特性多項式（というか、最小多項式の倍元）を使って多項式除算を行う、ということだった。この方法が適用できるのは、行列の特性多項式が既知の場合である。特に線型漸化式に付随する行列は特性多項式が既知なので、この方法を適用できる。
この問題について軽くググった感じ、日本の競プロ界隈ではこのアルゴリズムはきたまさ (kitamasa) 法と呼ばれているようだ。いちいち名前をつけるの好きだね……（名前があった方が解説に便利なのはわかる）。
前回の記事に登場したFibPairとは何だったのか
「行列の特性多項式を使ってべき乗を計算する」という観点から前回の記事「最速のフィボナッチ数計算を考える」に登場したFibPairというモノイド\begin{gather*}\mathsf{FibPair}=(\mathbf{Z}\times\mathbf{Z},\,\cdot\,)\\(a,b)\cdot(c,d):=(ad+bc-ac,ac+bd)\end{gather*}を考えると、これが何だったのか、そしてこれがなぜモノイドとなるのかはもはや明らかである。
そう、FibPair とは多項式の剰余環 \(\mathbf{Z}[X]/(X^2-X-1)\) （の積に関するモノイド）に他ならないのだ。実際、この剰余環で積を考えると\[(-aX+b)(-cX+d)\equiv -(ad+bc-ac)X+ac+bd\pmod{X^2-X-1}\]となり、FibPairと一致する。（不自然にマイナスがついているのは、FibPairの定義が不自然だったからだと考えられる。）
前回の記事の時点でこの視点に至れなかった筆者の不明を恥じる次第である。FibPairがモノイドの法則を満たすことを手計算でチェックしたのがアホらしい……。
ピンバック: 最速のフィボナッチ数計算を考える | 雑記帳
T _ { K } \in \left\langle I , T _ { K } , T _ { K } ^ { 2 } , \ldots , T _ { K } ^ { K – 1 }\right\rangle これ N 乗抜けているよ。
修正しました。ありがとうございます
ピンバック: Haskellerのためのモノイド完全ガイド | 雑記帳
ピンバック: AtCoder Beginners Contest 122 の D の別解 (ABC122-D) | 雑記帳