ferinの競プロ帳

競プロについてのメモ

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;
}