kmjp's blog

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

Codeforces #460 Div2 E. Congruence Equation

Eまでは調子よかったのにFでグダって残念。
http://codeforces.com/contest/919/problem/E

問題

正整数A,B,Xと素数Pが与えられる。
1~Xの間の整数Nのうち、N*A^N = B (mod P)を満たすものはいくつあるか。

解法

Nの部分はmod Pでは周期P、A^Nの部分はmod Pでは周期P-1である。
よってN*A^N全体では周期はP(P-1)である。

A^Nが周期P-1なので、ここでNを0~(P-2)まで総当たりしよう。
するとN mod (P-1)が決められた状態でN=B/(A^N) (mod P)がわかるので、中国人剰余定理の要領でNの候補がわかる。
あとはN+P(P-1)kの形の整数も解なので、X以下の範囲でいくつこの形の式があるか求めればよい。

ll A,B,P,X;
ll mo;
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;
}

ll ext_gcd(ll p,ll q,ll& x, ll& y) { // get px+qy=gcd(p,q)
	if(q==0) return x=1,y=0,p;
	ll g=ext_gcd(q,p%q,y,x);
	y-=p/q*x;
	return g;
}

pair<ll,ll> crt(ll a1,ll mo1,ll a2,ll mo2) { // return (x,y) y=lcm(a1,a2),x%mo1=a1,x%mo2=a2
	ll g,x,y,z;
	g=ext_gcd(mo1,mo2,x,y);
	a1=(a1%mo1+mo1)%mo1;a2=(a2%mo2+mo2)%mo2;
	if(a1%g != a2%g) return pair<ll,ll>(-1,0); // N/A
	ll lcm=mo1*(mo2/g);
	if(lcm<mo1) return pair<ll,ll>(-2,0); // overflow
	ll v=a1+((a2-a1)%lcm+lcm)*x%lcm*(mo1/g);
	return make_pair(((v%lcm)+lcm) % lcm,lcm);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>A>>B>>P>>X;
	mo=P;
	ll NA=1;
	ll ret=0;
	for(ll N=0;N<=P-2;N++) {
		ll R=B*modpow(NA)%mo;
		
		pair<ll,ll> p=crt(N,P-1,R,P);
		if(p.first<=X) {
			ret+=1+(X-p.first)/p.second;
		}
		NA=NA*A%mo;
	}
	cout<<ret<<endl;
}

まとめ

これは割とサクサク思いつけた。