kmjp's blog

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

TopCoder SRM 704 Div1 Medium ModEquation

手こずったけど解けてよかった。
https://community.topcoder.com/stat?c=problem_statement&pm=14469

問題

整数N,K,vが与えられる。
各要素0~(K-1)のいずれかの整数をとるN要素の数列X[i]のうち、全要素の積をとったとき、Kで割った余りがvとなるのは何通りあるか。

Div2 Hardと違い、

  • Nは50以下と小さい
  • Kは10^9以下と大きい
  • vは最大1000個与えられ、それぞれに対し答える必要がある

解法

まずはvは置いておいて、N個0~(K-1)の値を掛けたとき何が起こるかを考えてみる。

Kの約数を列挙し、昇順にD[i]とする。
0~(K-1)のうち、KとのGCDがD[i]となるようなものは複数(C[i]個)ある。
たとえば、K=72に対しKとのGCDが4になる値は4,20,28,44,52,68と6つある。

対称性よりX[i]を掛けていくとこれらの登場頻度は同じと推測できる。
だとすると、逆にX[i]を掛けていってKとのGCDがちょうどD[i]となる組み合わせ数が求まったとき、それらはC[i]個分の値になるケースを足し合わせたものなので、逆にC[i]で割れば個々の値を求めることができる。

よって、まずはvの値は置いておいて、X[i]をN個掛けたとき、途中過程でKとのGCDがD[i]であるようなものが何個あるかをDPしていくとよい。
この作業も上記C[i]の値を用いると、0~(K-1)まで掛け合わせるパターンをK通り試す必要はなく、0~(K-1)をそれぞれKとのGCDで分類することで、結局Kの約数個分の状態のみDPすればよくなる。
10^9以下の最大の高度合成数は735134400で約数が1344個のため、このDPは50*1344*1344程度のループで済む。
(以下のコードではもう少し削減している)

あとは各vに対しGCD(K,v)をもとに上記DPの結果を参照するだけ。

vector<ll> enumdiv(ll n) {
	vector<ll> S;
	for(ll i=1;i*i<=n;i++) if(n%i==0) {S.push_back(i); if(i*i!=n) S.push_back(n/i); }
	sort(S.begin(),S.end());
	return S;
}

int cnt[2020][2020];
vector<int> tar[2020];
ll from[2020], to[2020];
vector<ll> ed;
ll mo=1000000007;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}


class ModEquation {
	public:
	vector <int> count(int n, int K, vector <int> query) {
		vector<int> R;
		ed=enumdiv(K);
		int D=ed.size();
		int i,j;
		
		FOR(i,D) {
			tar[i].clear();
			for(j=i;j<D;j++) if(ed[j]%ed[i]==0) tar[i].push_back(j);
		}
		
		ZERO(cnt);
		for(i=ed.size()-1;i>=0;i--) {
			cnt[0][i]=K/ed[i];
			for(j=i+1;j<ed.size();j++) if(ed[j]%ed[i]==0) cnt[0][i]-=cnt[0][j];
		}
		for(i=1;i<D;i++) {
			FOR(j,D) cnt[i][lower_bound(ALL(ed),__gcd(ed[j]*ed[i],(ll)K))-ed.begin()]+=cnt[0][j];
		}
		
		ZERO(from);
		from[0]=1;
		while(n--) {
			ZERO(to);
			FOR(i,D) FORR(r,tar[i]) (to[r] += from[i]*cnt[i][r])%=mo;
			swap(from,to);
		}
		
		FORR(q,query) {
			int id=lower_bound(ALL(ed),__gcd(q,K))-ed.begin();
			R.push_back(from[id]*modpow(cnt[0][id])%mo);
		}
		return R;
	}
}

まとめ

これあんまりテストせずに先にSubmitしてしまったのだけど、その後K=72とか120とか約数の多そうな所で総当たりの結果と一致したので、Systest時にはかなり自信を持てた。