kmjp's blog

競技プログラミング参加記です

Codeforces #251 Div2 E. Devu and Birthday Celebration

ちょっと手こずったけど本番に解けてよかった。
http://codeforces.com/contest/439/problem/E

問題

以下のクエリQ個を処理せよ。

各クエリは2つの自然数N,Fからなる。
F要素からなる自然数の数列のうち、要素の合計がNで最大公約数が1となるものの組み合わせ数を求めよ。

解法

最大公約数の縛りを抜いて考えると、自然数NをF個に分ける分け方f(N,F)は重複組み合わせの考え方より(各要素最低1必要なことを考えると) f(N,F) = {}_{N-F} H_{F} = {}_{N-1} C_{F-1}である。

Nが異なる素因数としてp1,p2,p3...を持っていたとすると、包除原理より求める答えはg(x)=f(x,F)として
 g(N) - (g(N/p1+g(N/p2)+g(N/p3)...) + (g(N/(p1*p2))+g(N/(p1*p3))...) - ...
とNを素因数奇数個で割った値に対する関数gの値を引いて、偶数個で割った値に対する関数gの値を足せばよい。
N<=10^5より、素因数の数は高々6個なので計算すべきf(x,F)は2^6=64通りで済む。

 {}_{N-1} C_{F-1}の計算を何度も行うので、事前に累積積を取っておき、CombinationをO(1)で計算できるようにしておこう。
計算量は各クエリで素因数分解するためO(Q*√N)。

int Q;
int N[100001],F[100001];
ll mo=1000000007;

//同一の要素を複数含まない素因数分解
vector<ll> enumdiv(ll n) {
	vector<ll> V;
	if(n==1) return vector<ll>(1,1);
	for(ll i=2;i*i<=n;i++) {
		if(n%i==0) V.push_back(i);
		while(n%i==0) n/=i;
	}
	if(n>1) V.push_back(n);
	return V;
}

ll combi(int N_, int C_) {
	const int NUM_=200001;
	static ll fact[NUM_+1],factr[NUM_+1];
	int i;
	if (fact[0]==0) {
		vector<ll> inv(NUM_ + 1);
		inv[1] = 1;
		for (i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		fact[0]=factr[0]=1;
		FOR(i,NUM_) fact[i+1]=fact[i]*(i+1)%mo, factr[i+1]=factr[i]*inv[i+1]%mo;
	}
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

void solve() {
	int f,i,j,k,l,x,y,mask;
	
	cin>>Q;
	FOR(i,Q) {
		cin>>N[i]>>F[i];
		
		if(N[i]==1) {
			cout<<1<<endl;
			continue;
		}
		ll ret=0;
		vector<ll> VV=enumdiv(N[i]);
		FOR(mask,1<<VV.size()) {
			x=N[i];
			y=1;
			FOR(j,VV.size()) if(mask&(1<<j)) x/=VV[j],y=-y;
			if(x<F[i]) continue;
			x-=F[i];
			ret += y*combi(x+F[i]-1,F[i]-1);
		}
		ret %= mo;
		if(ret<0) ret+=mo;
		cout << ret << endl;
	}
}

まとめ

最初CombinationをO(N)で計算してしまいTLEした。
本番中に事前計算によりO(1)で計算するライブラリを作り何とか解けた。