선형점화식 빠르게 계산하기

$D_0, D_1, \cdots, D_{k-1}$과 $c_1, c_2, \cdots, c_k$가 주어졌을 때, $i \ge k$인 $D_i$를 다음과 같은 선형점화식으로 구할 수 있다고 하자.

$$D_i = \sum_{j=1}^{k}{c_jD_{i-j}} = c_1D_{i-1} + c_2D_{i-2} + \cdots + c_kD_{i-k}$$

이러한 선형점화식을 가지는 가장 유명한 예시는 피보나치 수열이다. 피보나치 수열은 위 식에서 $k=2$, $D_0 = 1$, $D_1 = 1$, $c_1 = 1$, $c_2 = 2$이다.

이러한 선형점화식이 주어졌을 때, $D_N$을 빠르게 구하는 방법을 알아보자.

 

1) 행렬 곱셈을 이용한 방법

$$\begin{pmatrix}
D_N \\
D_{N-1}\\
D_{N-2}\\
\vdots \\
D_{N-k+1}
\end{pmatrix}
=
\begin{pmatrix}
c_1 & c_2 & \cdots & c_{k-1} & c_{k} \\
1 & 0 & \cdots & 0 & 0 \\
0 & 1 & \cdots & 0 & 0\\
\vdots & \vdots & \ddots & \vdots & \vdots  \\
0 & 0 & \cdots & 1 & 0
\end{pmatrix}
\begin{pmatrix}
D_{N-1} \\ D_{N-2} \\ D_{N-3} \\ \vdots \\ D_{N-k}
\end{pmatrix}$$

주어진 선형점화식을 위와 같이 행렬식으로 표현할 수 있다. 이 $k \times k$ 크기의 행렬을 $A$라고 하자. 이는 다음과 같이 거듭제곱을 이용하여 표현할 수 있다.

$$\begin{pmatrix}
D_N \\
D_{N-1}\\
D_{N-2}\\
\vdots \\
D_{N-k+1}
\end{pmatrix}
=
\begin{pmatrix}
c_1 & c_2 & \cdots & c_{k-1} & c_{k} \\
1 & 0 & \cdots & 0 & 0 \\
0 & 1 & \cdots & 0 & 0\\
\vdots & \vdots & \ddots & \vdots & \vdots  \\
0 & 0 & \cdots & 1 & 0
\end{pmatrix}^{N-k+1}
\begin{pmatrix}
D_{k-1} \\ D_{k-2} \\ D_{k-3} \\ \vdots \\ D_{0}
\end{pmatrix} = A^{N-k+1}
\begin{pmatrix}
D_{k-1} \\ D_{k-2} \\ D_{k-3} \\ \vdots \\ D_{0}
\end{pmatrix}$$

즉, 정수 $a$의 $a^N$을 분할정복을 통해 $O(\lg N)$ 시간복잡도에 구할 수 있다. 이와 같은 방법으로 크기 $k \times k$인 행렬 $A$의 $A^N$을 $O(k^3 \lg N)$ 시간복잡도에 구할 수 있다. 이렇게 $A^{N-k+1}$을 구하면 $D_N$을 계산할 수 있다.

 

2) 키타마사법

만약 $k$가 크다면 $O(k^3 \lg n)$ 시간복잡도는 느릴 것이다. 좀 더 빠르게 구하는 방법을 예시를 통해 알아보자.

$k=2$이고, 점화식은 $D_i = 2D_{i-1}+3D_{i-2}$이다. $D_5$를 구하는 과정을 살펴보자. 이제 $D_5 = aD_1 + bD_0$ 꼴로 표현해보자.

$$D_5 = 2D_4+3D_3 = 2(2D_3+3D_2)+3D_3\\ = 7D_3+6D_2 = 7(2D_2+3D_1)+6D_2\\ = 20D_2 + 21D_1 = 20(2D_1+3D_0)+21D_1\\ = 61D_1+60D_0$$

$D_1$과 $D_0$을 알고 있기 때문에 $D_5$를 바로 계산할 수 있다.

 

점화식을 다항식으로 표현해보자. $x^i = 2x^{i-1}+3x^{i-2}$이며, $x^2 = 2x + 3$, 즉, $x^2-2x-3 = 0$이다. 다항식 $x^N$를 $x^2-2x-3$으로 나눠보면 다음과 같이 표현할 수 있다.

$$x^N = Q(x^2-2x-3) + R$$

다항식 $Q$는 나눗셈의 몫이 되며, 다항식 $R$은 나머지가 된다. 그리고 다항식 $R$의 차수는 2보다 작다. 만약, 다항식 $x^5$을 $x^2-2x-3$으로 나눠보면 다음과 같이 표현된다.

$$x^5 = (x^3+2x^2+7x+20)(x^2-2x-3) + 61x + 60$$

즉, 다항식 $R$의 계수가 $D_N = aD_1 + bD_0$ 꼴에서 $a$와 $b$가 된다.

 

일반화하면 점화식을 $g(x) = x^k - c_1x^{k-1} - \cdots - c_{k-1}x - c_k = 0$으로 나타내고 $x^N$을 $g(x)$로 나눈 나머지를 구하면 $D_N$을 구할 수 있다.

$$x^N = 
\begin{cases}
x^{\lfloor\frac{N}{2}\rfloor}x^{\lfloor\frac{N}{2}\rfloor}x & \mod g(x), & \text{if } N \text{ is odd}\\
x^{\frac{N}{2}}x^{\frac{N}{2}} & \mod g(x), & \text{otherwise}
\end{cases}$$

다항식의 경우에도 위 식이 성립하므로 $x^N$을 $g(x)$로 나눈 나머지를 $O(\lg N)$번의 곱셈과 나눗셈으로 구할 수 있다. 매 순간 $g(x)$로 나눈 나머지만 가지고 있으면 되므로, 곱하고 나누는 다항식의 차수는 $k$ 보다 작다.

 

차수가 $k$인 다항식의 곱셈과 나눗셈을 $O(k^2)$ 시간복잡도에 할 수 있는데, 이 경우 전체 시간복잡도는 $O(k^2 \lg N)$이 된다. 1번 방법과 비교했을 때 $k$ 배 빠른 방법이다.

 

3) 키타마사법 + NTT + 빠른 다항식 나눗셈

2번 방법에서 다항식의 곱셈과 다항식의 나눗셈에 $O(k^2)$ 시간이 필요했으므로 전체 시간복잡도가 $O(k^2 \lg N)$이었다.

 

만약, 문제의 MOD 값이 $a2^b+1$ 꼴의 소수이며 원시근을 알고 $k < 2^b$라면 NTT가 가능하다. 다항식의 곱셈은 NTT 그 자체로 $O(k \lg k)$ 시간에 가능하며, 나눗셈 또한 이 글에 있는 방법을 통해 $O(k \lg k)$에 가능하다.

 

전체 시간복잡도는 $O(k \lg k \lg N)$이 된다.

 

4) Bostan Mori 알고리즘

Bostan Mori 알고리즘에 대한 자세한 설명은 이 글을 참고하자.

 

이 알고리즘 또한 NTT를 필요로 하므로 3번 방법과 사용 가능한 조건이 동일하다. 시간복잡도 또한 $O(k \lg k \lg N)$으로 3번 방법과 동일하다. 그러나 다항식의 나눗셈을 필요로 하지 않는다는 장점이 있다. 다항식의 나눗셈의 경우, 시간복잡도는 $O(k \lg k)$이지만, 상수가 크고 별도로 구현해야 한다는 번거로움이 있기 때문에 이 알고리즘은 구현이 상대적으로 편하며 실행 시간이 3번 방법보다 빠르다.

 

*) 예시 문제

1번 방법의 기본 연습 문제다. 

2749번: 피보나치 수 3

MOD 값이 $1999 = 999\times2^1+1$이기 때문에 NTT를 쓸 수 없다. 2번 방법, $O(k^2 \lg N)$으로 구현해야 한다.

15572번: 블록 4

MOD 값이 $104857601 = 25 \times 2^{22}+1$ 로 NTT를 쓸 수 있다. 원시근은 2부터 차례대로 구해보면 되는데, 3이 원시근이 된다. $k$의 최댓값이 30000으로 크기 때문에 1번, 2번 방법으로 풀 수 없다.

13725번: RNG

아래 코드는 BOJ 13725 RNG 문제를 해결하는 코드이며, mod int와 NTT를 이용한 다항식 템플릿을 포함한 코드다.

#include <bits/stdc++.h> using namespace std; using lld = long long; template <int MOD> struct MINT{ int v{}; MINT() = default; constexpr MINT(lld v_){ v = -MOD <= v_ && v_ < MOD ? v_ : v_ % MOD; if (v < 0) v += MOD; } friend istream& operator >> (istream& is, MINT& a){ lld t; is >> t; a = MINT(t); return is; } friend ostream& operator << (ostream& os, const MINT& a){ return os << a.v; } bool operator == (const MINT& ot)const{ return v == ot.v; } bool operator != (const MINT& ot)const{ return v != ot.v; } static constexpr MINT pow(MINT a, lld b){ MINT ret = 1; for (;b;b>>=1,a*=a) if (b&1) ret *= a; return ret; } static constexpr MINT inv(const MINT& t){ return pow(t, MOD-2); } constexpr MINT operator - ()const{ return MINT(-v); } constexpr MINT& operator += (const MINT& ot){ if ((v += ot.v) >= MOD) v -= MOD; return *this; } constexpr MINT& operator -= (const MINT& ot){ if ((v -= ot.v) < 0) v += MOD; return *this; } constexpr MINT& operator *= (const MINT& ot){ v = (lld)v*ot.v%MOD; return *this; } constexpr MINT& operator /= (const MINT& ot){ *this *= inv(ot); return *this; } constexpr MINT operator + (const MINT& ot)const{ return MINT(*this) += ot; } constexpr MINT operator - (const MINT& ot)const{ return MINT(*this) -= ot; } constexpr MINT operator * (const MINT& ot)const{ return MINT(*this) *= ot; } constexpr MINT operator / (const MINT& ot)const{ return MINT(*this) /= ot; } constexpr operator int32_t ()const{ return v; } constexpr operator int64_t ()const{ return v; } }; template <int prime = (119<<23)+1, int primitive_root = 3> struct FFT{ using T = MINT<prime>; static constexpr int P = prime; static constexpr int PR = primitive_root; static constexpr int PRI = T::inv(PR); static constexpr int MSZ = (P - 1) & (1 - P); static constexpr int W = T::pow(PR, (P - 1) / MSZ); static constexpr int WI = T::pow(PRI, (P - 1) / MSZ); static void fft(int sz, vector<T>& a, bool inv){ T w = T::pow(inv ? WI : W, MSZ / sz); for (int ssz=sz>>1;ssz;ssz>>=1){ for (int i=0;i<sz;i+=ssz<<1){ T wt = 1; for (int j=0;j<ssz;++j){ T& lft = a[i + j]; T& rgh = a[i + j + ssz]; T lftold = lft; lft += rgh; rgh = (lftold - rgh) * wt; wt *= w; } } w *= w; } if (inv){ T invsz = T::inv(sz); for (int i=0;i<sz;i++) a[i] *= invsz; } for (int j=0,i=1;i<sz;i++){ int dg = sz>>1; for (;j&dg;dg>>=1) j ^= dg; j ^= dg; if (i < j) swap(a[i], a[j]); } } static vector<T> convolution(const vector<T>& a, const vector<T>& b){ int sz = a.size() + b.size(); int fftsz = sz; while (fftsz & (fftsz - 1)) fftsz += fftsz & -fftsz; vector<T> av(fftsz); vector<T> bv(fftsz); for (int i=0,ilim=size(a);i<ilim;i++) av[i] = a[i]; for (int i=0,ilim=size(b);i<ilim;i++) bv[i] = b[i]; fft(fftsz, av, false); fft(fftsz, bv, false); for (int i=0;i<fftsz;i++) av[i] *= bv[i]; fft(fftsz, av, true); return av; } }; // alternatives using FFT1 = FFT<1'012'924'417, 5>; // (483<<21)+1 using FFT2 = FFT<1'004'535'809, 3>; // (479<<21)+1 template <int prime = (119<<23)+1, int primitive_root = 3> struct Poly{ using T = MINT<prime>; using fft_t = FFT<prime, primitive_root>; vector<T> a; Poly() = default; Poly(T a0): a(1, a0){ normalize(); } Poly(const vector<T>& a): a(a) { normalize(); } int size()const{ return a.size(); } int deg()const{ return a.size()-1; } T& operator [](int i){ return a[i]; } typename vector<T>::const_iterator begin()const{ return a.begin(); } typename vector<T>::const_iterator end()const{ return a.end(); } void push_back(const T& v){ a.push_back(v); } template<typename... Args> void emplace_back(Args&&... args){ a.emplace_back(args...); } void pop_back(){ a.pop_back(); } void resize(int n){ a.resize(n); } void normalize(){ while (!a.empty() && a.back() == T(0)) a.pop_back(); } Poly reversed()const{ auto b = a; reverse(b.begin(), b.end()); return b; } Poly trim(int n)const{ return vector<T>(a.begin(), a.begin() + min(n, size())); } Poly inv(int n){ Poly q(T(1) / a[0]); for (int i=1;i<n;i<<=1){ Poly p = Poly(2) - q*trim(i<<1); q = (p * q).trim(i<<1); } return q.trim(n); } Poly& operator *= (const T& x){ for (T& v: a) v *= x; normalize(); return *this; } Poly& operator /= (const T& x){ return *this *= T(1) / x; } Poly& operator += (const Poly& ot){ if (size() < ot.size()) a.resize(ot.size()); for (int i=0;i<ot.size();i++) a[i] += ot.a[i]; normalize(); return *this; } Poly& operator -= (const Poly& ot){ if (size() < ot.size()) a.resize(ot.size()); for (int i=0;i<ot.size();i++) a[i] -= ot.a[i]; normalize(); return *this; } Poly& operator *= (const Poly& ot){ *this = fft_t::convolution(a, ot.a); normalize(); return *this; } Poly& operator /= (const Poly& ot){ if (deg() < ot.deg()) return *this = Poly(); int sz = deg()-ot.deg()+1; Poly ra = reversed().trim(sz), rb = ot.reversed().trim(sz).inv(sz); *this = (ra*rb).trim(sz); while (size() < sz) emplace_back(); reverse(a.begin(), a.end()); normalize(); return *this; } Poly& operator %= (const Poly& ot){ if (deg() < ot.deg()) return *this; Poly tmp = *this; tmp /= ot; tmp *= ot; *this -= tmp; normalize(); return *this; } Poly operator * (const T& x)const{ return Poly(*this) *= x; } Poly operator / (const T& x)const{ return Poly(*this) /= x; } Poly operator + (const Poly& ot)const{ return Poly(*this) += ot; } Poly operator - (const Poly& ot)const{ return Poly(*this) -= ot; } Poly operator * (const Poly& ot)const{ return Poly(*this) *= ot; } Poly operator / (const Poly& ot)const{ return Poly(*this) /= ot; } Poly operator % (const Poly& ot)const{ return Poly(*this) %= ot; } }; constexpr int MOD = 104857601, PR = 3; // (25 << 22) + 1 using mint_t = MINT<MOD>; using poly_t = Poly<MOD, PR>; int main(){ cin.tie(0)->sync_with_stdio(0); int K; lld N; cin >> K >> N; vector<mint_t> A(K), B(K); for (auto& v: A) cin >> v; for (auto& v: B) cin >> v, v *= -1; reverse(B.begin(), B.end()); B.emplace_back(1); N--; if (N < K) return cout << A[N] << '\n', 0; poly_t res(1), v({0, 1}); for (int i=0;i<60;i++,v=v*v%B) if (N>>i&1) res = res*v%B; mint_t ans = 0; for (int i=0;i<K;i++) ans += res[i]*A[i]; cout << ans << '\n'; }