kmjp's blog

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

AtCoder ABC #270 (トヨタ自動車プログラミングコンテスト2022) : G - Sequence in mod P

ARC

ギリギリ全完できてよかった。
https://atcoder.jp/contests/abc270/tasks/abc270_g

問題

素数Pと非負整数S,G,A,Bが与えられる。
整数列Xを以下のように取る。

  • X[0]=S
  • X[i+1]=(A*X[i]+B) % P

X[i]=Gとなるiが存在すれば、最小のものを求めよ。

解法

コーナーケースを片付けておく。

  • S=Gなら、X[0]=G
  • A=0なら、X={S,B,B,B...}となるので、解は0か1か存在しないのいずれか。
  • A=1なら、(S+iB)≡G (mod P)となるiを求めればよいのでi≡(G-S)/Bで計算できる。

以後Aが2以上のケースを考える。
\displaystyle X_n = (A^nS + B \times \sum_{i=0}^{n-1} A^i) % P
なので、式変形すると以後(mod P)を省略して
\displaystyle A^nS + B \times \frac{A^n-1}{A-1} \equiv G
\displaystyle (S(A-1)+B)A^n \equiv G(A-1)+B
\displaystyle A^n \equiv \frac{G(A-1)+B}{S(A-1)+B}
とすると、これは離散対数方程式になるのでBaby-Step Giant-Stepで解くことができる。

int T;
ll A,B,S,G,mo;

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

map<int,int> M;
ll modlog(int g,int a) {  // return g^x=a mod mo
	M.clear();
	ll cur=1;
	int sq=sqrt(mo);
	int i;
	FOR(i,sq) {
		if(M.count(cur)==0) M[cur]=i;
		cur=cur*g%mo;
	}
	
	ll step=cur;
	cur=1;
	ll num=0;
	int lp=0;
	while(1) {
		ll x=a*modpow(cur)%mo;
		if(lp++>sq+5) return 1LL<<50;
		if(M.count(x)) return num+M[x];
		cur=cur*step%mo;
		num+=sq;
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>mo>>A>>B>>S>>G;
		if(S==G) {
			cout<<0<<endl;
			continue;
		}
		if(A==0) {
			if(B==G) {
				cout<<1<<endl;
			}
			else {
				cout<<-1<<endl;
			}
			continue;
		}
		if(A==1) {
			if(B==0) {
				cout<<-1<<endl;
			}
			else {
				ll ret=(G-S+mo)*modpow(B)%mo;
				cout<<ret<<endl;
			}
			continue;
		}
		ll L=((A-1)*S+B)%mo;
		ll R=(G*(A-1)+B)%mo;
		if(L==0||R==0) {
			cout<<-1<<endl;
			continue;
		}
		R=R*modpow(L)%mo;
		
		ll ret=modlog(A,R);
		if(ret==1LL<<50) {
			cout<<-1<<endl;
		}
		else {
			cout<<ret<<endl;
		}
	}
}

まとめ

本番、BSGSのライブラリがバグっててタイムロスした…。