AtCoder ABC #245 : Ex - Product Modulo 2

本番出てたらコーナーケースで落としてそう。
https://atcoder.jp/contests/abc245/tasks/abc245_h

問題

整数K,N,Mが与えられる。
以下を満たす数列Aは何通りか。

• AはK要素で、それぞれ0以上M未満の整数値を取る。
• Aの各要素の積をMで割った余りはNに等しい。

解法

Mを素因数分解して(p1^q1)*(p2^q2)....と表せたとする。
Aを0以上pi^qi未満の値を取るケースを考え、prod(A)=N (mod (pi^qi))となるものを各(pi,qi)毎に求めてその積を取ればよい。

以下、1つの(p, q)に対しprod(A)=N (mod (p^q))となるAを数え上げることを考えよう。
各要素は0～(p^q-1)を等確率に取るので、prod(A)の値のうち、pで割れる回数が等しいものは登場回数も等しい。
そこでそのようなものはまとめて扱おう。
B(n,i) := Aのprefix n要素の積のうち、p^iで割り切れるがp^(i+1)で割り切れないものの数
とすると、Bの要素数は高々q+1個なので、行列累乗で高速にB(K,*)を求められる。

N=r*p^iと表現できるとき、B(K,i) / (p^iで割り切れるがp^(i+1)で割り切れないものの数)を求めればよい。
分母が998244353の倍数となる可能性があるので、B(K,i)を求める過程ではmod (998244353^2)で計算すると良い。

int K;
ll N,M;
const ll mo=998244353;
const __int128 mo2=(__int128)mo*mo;
ll num[44];
ll P[44];

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;
}


const int MAT=41;
struct Mat { __int128 v[MAT][MAT]; Mat(){ZERO(v);};};

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo2;
	return r;
}

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;
}

Mat powmat(ll p,Mat a,int n=MAT) {
	int i,x,y; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(i,n) r.v[i][i]=1;
	while(p) {
		if(p%2) r=mulmat(r,a,n);
		a=mulmat(a,a,n);
		p>>=1;
	}
	return r;
}

ll hoge(ll p,int q) {
	ll M=1;

	int i,j;
	P[0]=1;
	FOR(i,q) P[i+1]=P[i]*p;
	num[q]=1;
	for(i=q-1;i>=0;i--) num[i]=P[q-i]-P[q-1-i];
	
	Mat A;
	FOR(i,q+1) FOR(j,q+1) A.v[min(i+j,q)][i]+=num[j];
	A=powmat(K,A,q+1);
	ll n=N%P[q];
	i=0;
	if(n==0) {
		i=q;
	}
	else {
		while(n%p==0) i++,n/=p;
	}
	
	if(num[i]%mo==0) {
		A.v[i][0]/=mo;
		num[i]/=mo;
	}
	ll a=A.v[i][0]%mo*modpow(num[i])%mo;
	return a;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>K>>N>>M;
	auto p=enumpr(M);
	ll ret=1;
	FORR2(a,b,p) (ret*=hoge(a,b))%=mo;
	cout<<ret<<endl;
}

まとめ

Mは998244353未満じゃダメだったんかな。