AtCoder ABC #335 (Sponsored by Mynavi) : G - Discrete Logarithm Problems

これは割とすんなり。
https://atcoder.jp/contests/abc335/tasks/abc335_g

問題

素数Pと、P未満の正整数からなるN要素の整数列Aが与えられる。
正整数kを用いてA[i]^k ≡ A[j] (mod P) となる(i,j)の対は何通りか。

解法

f(v)を、v^k≡1 (mod P)となる最小のkとすると、f(v)は(P-1)の約数なので、(P-1)を(P-1)の素因数で割って行ってその値を求めよう。
A[i]^k ≡ A[j] (mod P)となる条件は、f(A[i])がf(A[j])の倍数の場合なので、f(A[i])をそれぞれ求めておけば容易に計算できる。

int N;
ll P,mo;

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

map<ll,int> enumpr(ll n) {
	map<ll,int> V;
	for(ll i=2;i*i<=n;i++) while(n%i==0) V[i]++,n/=i;
	if(n>1) V[n]++;
	return V;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	cin>>N>>mo;
	map<ll,int> cand=enumpr(mo-1);
	
	map<ll,ll> M;
	FOR(i,N) {
		ll num=mo-1;
		ll A;
		cin>>A;
		FORR2(a,b,cand) {
			FOR(x,b) {
				if(modpow(A,num/a)>1) break;
				num/=a;
			}
		}
		M[num]++;
	}
	
	ll ret=0;
	FORR2(a1,b1,M) FORR2(a2,b2,M) if(a1%a2==0) ret+=b1*b2;
	
	cout<<ret<<endl;
	
	
	
}

問題

P^2が64bit超える点でちょっとハマってしまった。