TopCoder SRM 449 Div1 Hard StairsColoring

問題

n段の階段があり、ちょうどn枚の長方形の板でできている。
長方形の板には重なりやはみだしがない。
このようにして作れる階段を全部つくり、階段を（一つの階段は一色で）それぞれk色のどれかで塗る。

このようにして出来上がるものは何通りあるか、mod 1000000123で求めよ。

制約条件

n, k≦10^9

方針

n段の階段の作り方について漸化式を立てる。
これは、大きい長方形を決めることが()を決めることに対応するので、
長さ2nの括弧からなる、括弧の対応する文字列の個数に等しい。
これはC(2n, n) / (n+1)!個である。

これをc(n)とすると、k^c(n)をmod 1000000123で求めたい。
1000000123は素数なので、フェルマーの小定理より、k^xは周期1000000122で循環する。

つまり、c(n)をmod 1000000122で求めればいいことになる。
1000000122を素因数分解して、それぞれの素因数についてのmodでc(n)を求めて、
中国剰余定理を使えばc(n)がmod 1000000122で求まる。

あとはk^xを繰り返し二乗法でmod 1000000123で求めればいい。

ソースコード

ll fact[10000];
ll extgcd(ll a, ll b, ll &x, ll &y){
	ll d = a;
	if(b){
		d = extgcd(b, a % b, y, x);
		y -= a/b * x;
	}
	else x = 1, y = 0;
	return d;
}
ll mod_inverse(ll a, ll m){
	ll x, y;
	extgcd(a, m, x, y);
	return (m + x%m) % m;
}
ll mod_fact(ll n, ll p, int &e){
	e = 0;
	if(n == 0) return 1;
	ll res = mod_fact(n/p, p, e);
	e += n/p;
	
	if(n/p % 2) return res * (p - fact[n%p]) % p;
	return res * fact[n%p] % p;
}
int calc(ll n, ll p){
	int e1, e2, e3;
	ll a1 = mod_fact(2*n, p, e1), a2 = mod_fact(n+1, p, e2), a3 = mod_fact(n, p, e3);
	if(e1 > e2 + e3) return 0;
	return a1 * mod_inverse(a2 * a3 % p, p) % p;
}

bool linear_congruence(const vector<ll> &A, const vector<ll> &B, const vector<ll> &M,
	ll &x, ll &m){
	x = 0; m = 1;
	
	rep(i,A.size()){
		ll a = A[i] * m, b = B[i] - A[i] * x, d = __gcd(M[i], a);
		if(b % d) return 0;
		ll t = b /d * mod_inverse(a/d, M[i]/d) % (M[i] / d);
		x = x + m * t;
		m *= M[i] / d;
	}
	x = (x + m) % m;
	return 1;
}

class StairsColoring {
  public:
  int coloringCount(int N, int K) {
		int ps[] = {2, 3, 11, 2089, 7253};
		vector<ll> a, b, m;
		rep(i,5){
			fact[0] = 1;
			rep(j, ps[i]) fact[j+1] = fact[j] * (j+1) % ps[i];
			b.pb(calc(N, ps[i]));
			m.pb(ps[i]);
			a.pb(1);
		}
		ll x, M;
		linear_congruence(a, b, m, x, M);
		dbg(x); dbg(M);
		rep(i, 5){
			cerr << a[i] << ", " << b[i] << ", " << m[i] << endl;
		}
		M++;
		ll ans = 1, y = K;
		for(; x; x /= 2){
			if(x & 1) ans = ans * y % M;
			y = y * y % M;
		}
		return ans;
  }
}: