Educational Codeforces Round 60 D. Magic Gems

問題ページ

解法

dp[i]=(i個分スペースを埋める方法が何通りあるか)としたDPを考える。i個スペースを埋めるためにはmagic gemを一つ置くパターンとmagic gemを分割しnormal gemをM個置くパターンがある。したがってこのDPの遷移は

• dp[i] = dp[i-m] + dp[i-1] (i >= m)
• dp[i] = 1 (i < m)

となる。N<=10^18なので普通にDPをすることはできないがこの漸化式は行列累乗を用いることN項目をO(M^3logN)で求めることができ解けた。

#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<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 ll MOD = 1000000007;

using vec = vector<ll>;
using mat = vector<vec>;

// A*Bを計算
mat mul(mat &A, mat &B) {
    mat C(A.size(), vec(B[0].size()));
    for(int i=0; i<(int)A.size(); ++i) {
        for(int k=0; k<(int)B.size(); ++k) {
            for(int j=0; j<(int)B[0].size(); ++j) {
                C[i][j] = (C[i][j] + A[i][k]*B[k][j]) % MOD;
            }
        }
    }
    return C;
}

// A*Bを計算
vec mul(mat &A, vec &B) {
    vec C(A.size());
    for(int i=0; i<(int)A.size(); ++i) {
        for(int j=0; j<(int)B.size(); ++j) {
            C[i] = (C[i] + A[i][j]*B[j]) % MOD;
        }
    }
    return C;
}

// A^nを計算 m*m行列のn乗はO(m^3logn)
mat pow(mat A, ll n) {
    mat B(A.size(), vec(A.size()));
    for(int i=0; i<(int)A.size(); ++i) B[i][i] = 1;
    while(n > 0) {
        if(n & 1) B = mul(B, A);
        A = mul(A, A);
        n >>= 1;
    }
    return B;
}

signed main(void) {
    cin.tie(0);
    ios::sync_with_stdio(false);

    ll n, m;
    cin >> n >> m;

    if(n == m) {
        cout << 2 << endl;
        return 0;
    }

    if(n < m) {
        cout << 1 << endl;
        return 0;
    }

    mat a(m, vec(m));
    a[0][0] = 1, a[0][m-1] = 1;
    REP(i, m-1) a[i+1][i] = 1;

    vec b(m, 1);
    b[0] = 2;

    a = pow(a, n-m);
    cout << mul(a, b)[0] << endl;

    return 0;
}