kmjp's blog

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

Codeforces #536 Div2 F. Lunar New Year and a Recursive Sequence

Codeforces

勉強になりました。
http://codeforces.com/contest/1106/problem/F

問題

ある数列をF(i)を考える。
F(1)～F(K-1)はいずれも1で、F(K)は不定とする。
P=998244353とするとき、i＞Kにおいて、 \displaystyle F(i) =F(i-1)^{B_1} \cdot F(i-2)^{B_2} \cdots F(i-K)^{B_K} mod P = \prod_{j=1}^K F(i-j)^{B_j} mod Pとする。
F(N)=Mであるとき、F(K)はいくつか。

解法

両辺logを取ると \displaystyle \log(F(i)) =\sum_{j=1}^K B_j \log(F(i-j)) \mod (P-1)である。
F(1)～F(K-1)は1ということはlog(F(1))～log(F(K-1)) = 0なので、いつもの漸化式を行列累乗を用いることで、log(F(K))を何倍するとlog(F(N))= log Mとなるかがわかる。

よって、仮にa * log(F(K)) = log(F(N)) = log(M)とすると、F(K)のa乗がMということになる。
そこでこの問題は、Mのa乗根を求めればよいことがわかった。

剰余環における冪根は、確率的な方法もあるようだがここでは別の解法で解くことにする。
まずmod Pにおける原始根gを求めよう。これはgを(P-1)の約数乗したとき、g^(P-1)以外では1にならないことを総当たりすればよい。
今回のPではg=3がそれにあたる。

求めたいのは x^a=Mだが、まず g^y=Mとなるyを求めよう。
元の問題は冪根だが、こちらは離散対数を求める問題なのでBaby-step Giant-Stepが利用できる。
yを分解し、 g^y = g^{a*b} = {g^a}^bとなるbが求められるなら、 x = g^bとなって解が求まる。
y,aがわかっているとき y = a*b \mod (P-1)となるbを求めたい。
これは２つ方法が考えられる。

  • 一つ目は拡張ユークリッドの互除法
  • 二つ目はaの逆数を求め y/a = b \mod (P-1)を求めること。
    • このためには、まずGCD(a,P-1)を求め、yがその倍数であることを確認しておこう。yがGCD(a,P-1)の倍数でないなら、aを何倍してもbにはならず、すなわち解なしである。
    • y', a', p' をy,a,(P-1)をGCD(a,P-1)で割ったものとする。a'のmod p'における逆数はa'をφ(p')-1乗すれば求められる。あとはb = y'/a' mod p'でbを求めればよい。
int K;
int B[101];
int N,M;
ll mo=998244353;

const int MAT=101;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	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]%=mo;
	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 modpow(ll a, ll n = mo-2, ll m=mo) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%m,a=a*a%m,n>>=1;
	return r;
}

int totient(int v) {
	int ret=v;
	for(int i=2;i*i<=v;i++) if(v%i==0) {
		ret=ret/i*(i-1);
		while(v%i==0) v/=i;
	}
	if(v>1) ret=ret/v*(v-1);
	return ret;
}

int mod_root(int p,int a) { // x^p=a mod mo
	vector<int> D;
	for(int i=2;i*i<=mo-1;i++) if((mo-1)%i==0) D.push_back(i),D.push_back((mo-1)/i);
	int g=2;
	while(1) {
		int ng=0;
		FORR(d,D) if(modpow(g,d)==1) ng=1;
		if(ng==0) break;
		g++;
	}
	
	ll cur=a;
	int rg=modpow(g);
	int mstep=sqrt(mo);
	map<int,int> M;
	int i;
	FOR(i,mstep+3) {
		M[cur]=i;
		cur=cur*rg%mo;
	}
	ll pg=modpow(g,mstep);
	int x=-1,step=0;
	cur=1;
	while(x==-1) {
		if(M.count(cur)) x=step+M[cur];
		M[cur]=step;
		cur=cur*pg%mo;
		step+=mstep;
	}
	// g^x=aなのでg^(p*q)=g^x=aとしてq=x/p (mod mo-1) mo-1は合成数なのでGCDで割って対応
	int tmo=mo-1;
	int gcd=__gcd(tmo,p);
	if(x%gcd) return -1;
	tmo/=gcd;
	x/=gcd;
	p/=gcd;
	return modpow(g,1LL*x*modpow(p,totient(tmo)-1,tmo)%tmo);
}


Mat A,C;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>K;
	FOR(i,K) {
		cin>>A.v[0][i];
		if(i) A.v[i][i-1]=1;
	}
	cin>>N>>M;
	
	mo--;
	C=powmat(N-K,A,K);
	mo++;
	
	// p**N=x;
	cout<<mod_root(C.v[0][0],M)<<endl;
	
	
}

まとめ

指数対数を行ったり来たりする問題。
勉強になったね。