SRM566 Div1 Medium(500) PenguinEmperor

PenguinEmperor

対称性より、街iから街jに到達する旅程の数と、街i+kから街j+kに到達する旅程の数は等しい。街の数をnとする。i日目とi+n日目は同じ動きをするので、街0から出発してn日後にそれぞれの街へ到達する旅程の数を覚えておけば、それを使って、n日後、2n日後…の旅程の数が求められる。このままではdaysPassed/n回の計算が必要だけど、バイナリ法を使うと、log₂(daysPassed/n)で計算できる。

#include <vector>
using namespace std;

const long long M =  1000000007;

vector<long long> mul( vector<long long> A, vector<long long> B )
{
    int n = (int)A.size();
    vector<long long> C(n);
    for ( int i=0; i<n/2+1; i++ )   //  Cは線対称
        for ( int j=0; j<n; j++ )
            C[i] += A[(i-j+n)%n]*B[j],
            C[i] %= M;
    for ( int i=n/2+1; i<n; i++ )
        C[i] = C[n-i];
    return C;
}

vector<long long> pow( vector<long long> A, long long p )
{
    int n = (int)A.size();
    vector<long long> B(n); B[0]=1LL;
    for ( ; p>0; p/=2 )
    {
        if ( p%2==1 )
            B = mul(B,A);
        A = mul(A,A);
    }
    return B;
}

class PenguinEmperor{public:
int countJourneys( int numCities, long long daysPassed )
{
    int n = numCities;

    //  i日目での移動
    vector<vector<long long> > A(n,vector<long long>(n));
    for ( int i=0; i<n; i++ )
        A[i][i] = A[i][(n-i)%n] = 1LL;

    //  numCities日目にi番目の街にいる旅程数
    vector<long long> B(n); B[0]=1LL;
    for ( int i=1; i<=n; i++ )
        B = mul(B,A[i%n]);

    vector<long long> C = pow(B,daysPassed/n);
    for ( long long i=daysPassed/n*n+1; i<=daysPassed; i++ )
        C = mul(C,A[i%n]);

    return (int)C[0];
}};