TDPC T - フィボナッチ
解法
きたまさ法メモ - よすぽの日記
この記事を自分なりに理解したメモ
問題では1-indexになっているが0-indexで扱うとする。つまり以下の数列Aの第n項を求めろという問題になる。
A[i] = 1 (i<k)
A[i] = A[i-1] + A[i-2] + … + A[i-k] (i>=k)
行列累乗を使えばO(K^3N)で解ける(cf. 蟻本p114)がこれでは当然TLEする。きたまさ法と呼ばれる線形K項間漸化式の第n項をO(K^2logN)で求めるアルゴリズムを用いる。関数f(m)を数列xを返す関数とする。数列xはA[m]=x[0]A[0]+x[1]A[1]+…+x[k-1]A[k-1]を満たすとする。f(n)を求めることができればA[n]も求まる。f(m)からf(m+1)、f(2m)を高速に求めることができればダブリング・繰り返し二乗法の要領で計算することでO(logn)回の計算でf(n)が求まる。
f(m)からf(m+1)は以下の要領でO(k)で求められる。
A[m] = x[0]A[0]+x[1]A[1]+…+x[k-1]A[k-1]
A[m+1] = x[0]A[1]+x[1]A[2]+…+x[k-1]A[k]
= x[0]A[1]+x[1]A[2]+…+x[k-1]*(A[0]+A[1]+…+A[k-1])
= x[k-1]A[0] + (x[0]+x[k-1])A[1] + … + (x[k-2]+x[k-1])A[k-1]
f(m)からf(2m)はO(k^2)で求められる。まずO(k^2)かけて上記と同じ要領でf(m),f(m+1),…,f(m+k-1)を求める。f(m)が返す数列のi番目をf(m,i)と書くことにするとA[2m]は以下のように書ける。
A[2m] = x[0]A[n]+x[1]A[n+1]+…+x[k-1]A[n+k-1]
= x[0]*(f(n,0)A[0] + f(n,1)A[1] + … + f(n,k-1)A[k-1]) + … + x[k-1]*(f(n+k-1,0)A[0] + …)
= (x[0]f(n,0) + x[1]f(n+1,0) + … + x[k-1]f(n+k-1,0))A[0] + (x[1]f(n,1) + …)A[1] + … + (x[k-1]f(n,k-1) + …)A[k-1]
以上のようにf(m)からf(m+1)、f(2m)を計算することがO(k^2)でできるため第n項をO(k^2logn)で求めることができた。
#include <bits/stdc++.h> using namespace std; using ll = long long; // #define int ll using PII = pair<ll, ll>; #define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i) #define REP(i, n) FOR(i, 0, n) #define ALL(x) x.begin(), x.end() template<typename T> T &chmin(T &a, const T &b) { return a = min(a, b); } template<typename T> T &chmax(T &a, const T &b) { return a = max(a, b); } template<typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; } template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); } template<typename T> vector<T> make_v(size_t a) { return vector<T>(a); } template<typename T,typename... Ts> auto make_v(size_t a,Ts... ts) { return vector<decltype(make_v<T>(ts...))>(a,make_v<T>(ts...)); } template<typename T,typename V> typename enable_if<is_class<T>::value==0>::type fill_v(T &t, const V &v) { t=v; } template<typename T,typename V> typename enable_if<is_class<T>::value!=0>::type fill_v(T &t, const V &v ) { for(auto &e:t) fill_v(e,v); } template<class S,class T> ostream &operator <<(ostream& out,const pair<S,T>& a){ out<<'('<<a.first<<','<<a.second<<')'; return out; } template<typename T> istream& operator >> (istream& is, vector<T>& vec){ for(T& x: vec) {is >> x;} return is; } template<class T> ostream &operator <<(ostream& out,const vector<T>& a){ out<<'['; for(T i: a) {out<<i<<',';} out<<']'; return out; } int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0}; // DRUL const int INF = 1<<30; const ll LLINF = 1LL<<60; const int MOD = 1000000007; // a_nをO(K^2logN)で求める vector<ll> dfs(vector<ll> a, vector<ll> d, ll n) { ll k = d.size(); if(n == k) return d; vector<ll> ret(k); if(n&1 || n<k*2) { auto v = dfs(a, d, n-1); ret[0] = v[k-1] * a[0] % MOD; FOR(i, 1, k) ret[i] = (v[i-1] + v[k-1] * a[i] % MOD) % MOD; } else { auto v = dfs(a, d, n/2); vector<vector<ll>> f(k, vector<ll>(k)); f[0] = v; FOR(i, 1, k) { f[i][0] = f[i-1][k-1] * a[0] % MOD; FOR(j, 1, k) f[i][j] = (f[i-1][j-1] + f[i-1][k-1] * a[j] % MOD) % MOD; } REP(i, k) REP(j, k) (ret[i] += v[j] * f[j][i] % MOD) %= MOD; } return ret; } ll kitamasa(vector<ll> a, vector<ll> d, ll n) { auto ret = dfs(a, d, n); ll ans = 0; REP(i, d.size()) (ans += d[i] * ret[i]) %= MOD; return ans; } signed main(void) { cin.tie(0); ios::sync_with_stdio(false); ll n, k; cin >> k >> n; n--; if(n < k) { cout << 1 << endl; return 0; } vector<ll> d(k, 1); cout << kitamasa(d, d, n) << endl; return 0; }