TopCoder SRM 566 Div1 Medium PenguinEmperor

問題

n個の都市が円上0, 1, 2, ..., n - 1, 0という順に並んでいる。
最初0の都市にいて、
1日目には1個隣の都市のどちらかに、
2日目には2個隣の都市のどちらかに、
...
とm日間移動を繰り返す。


m日の移動が終わった後に都市0にいるような移動ルートは何通りあるか、
mod 10^9 + 7で求めよ。
ただし、二つのルートは、同じ日に違う都市にいるような日がひとつでもあるとき異なるものとして数える。

制約条件

n≦350
m≦10^18

方針

行列累乗。
n + 1日後は1日目と同じ移動方法しかないので、
M[i][j]をiからjまでを1〜n日までで移動する場合の数とする。


すると、mのうち、nで割り切れる部分の移動方法はM^(m / n)の行列になり、繰り返し二乗法で計算でき、
余りの部分は350以下なので、愚直に遷移行列をかければいい。


ただしそのままやると時間計算量O(n^4 log m)でどうみてもTLEで、
遷移行列が全部巡回行列であることを利用して、一行分だけを計算することで、
O(n^3 log m)の時間計算量で計算する必要がある。

ソースコード

inline vi operator*(const vi & a, const vi &b){
	int n = a.size();
	vi res(n);
	rep(i, n) rep(j, n){
		res[i] += (ll)a[j] * b[(j - i + n) % n] % mod;
		res[i] %= mod;
	}
	return res;
}
inline vi pow(vi a, ll m){
	vi res(a.size());
	res[0] = 1;
	for(; m; m /= 2){
		if(m & 1) res = res * a;
		a = a * a;
	}
	return res;
}

//遷移行列を作る関数
vi gen(int n, int m){
	vi res(n);
	res[m % n]++;
	if(res[(n - m % n) % n] == 0) res[(n - m % n) % n]++;
	return res;
}

class PenguinEmperor {
	public:
	int countJourneys(int n, long long d) {
		
		ll a = d / n;
		int b = d % n;
		
		vi res(n);
		res[0] = 1;
		//余りの部分
		rep(i, b) res = res * gen(n, i + 1);

		割り切れる部分
		vi pw = res;
		for(int i = b; i < n; i++) pw = pw * gen(n, i + 1);
		res = res * pow(pw, a);
		
		return res[0];
	}
};