Codeforces #281 Div2 E. Vasya and Polynomial

難しいようで、丁寧に解けばさほどでもない。
http://codeforces.com/contest/493/problem/E

問題

xのi乗の項の係数a[i]が非負整数である多項式P(x)を考える。
N,A,Bが与えられるので、P(N)=A、P(P(N))=Bとなる多項式P(x)がいくつあるか答えよ。

解法

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

• N＞Aの場合、A==BならP(x)=A一択である。A!=Bなら解はない。
• N==Aの場合、
  • A!=Bなら解はない。
  • N==A==B==1の場合、A[0]+A[1]+A[2]...=1となる答えは無限にある。
  • N==A==B＞1の場合、P(x)=A=xの２つの買いがある。

残るケースはN＜Aの場合である。
まず、P(N)=A、P(A)=Bより、N^m≦A、A^m≦Bとなるmを先に求めることで、あり得る最高次の次数mを求めておく。

あとは上位の項からDFSで係数A[i]を求めていけば良い。
途中で枝刈りをすることで、メモ化しなくても高速に終わる。

ll N,A,B;
ll mo=1000000007;
ll PP[2][65];
int mm;

ll dpdp(ll L,ll R,int m) {
	if(m==0) return L==R;
	
	ll ret=0,a;
	if(m==1) {
		if((R-L)%(A-N)) return 0;
		ll a1=(R-L)/(A-N);
		ll a0=A-a1*N;
		return a0>=0 && a1>=0;
	}
	
	FOR(a,min(L/PP[0][m],R/PP[1][m])+1) {
		ll L2=L-a*PP[0][m];
		ll R2=R-a*PP[1][m];
		if(L2<0 || R2<0 || L2>R2) break;
		if(L2==0) {
			ret += L2==R2;
		}
		else {
			if(R2/L2 > PP[1][m]/PP[0][m]) continue;
			ret += dpdp(L2,R2,m-1);
		}
		ret %= mo;
	}
	return ret;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>A>>B;
	
	if(N==A) {
		if(A!=B) return _P("0\n");
		if(N==1) return _P("inf\n");
		return _P("2\n");
	}
	if(N>A) return _P("%d\n",A==B);
	
	PP[0][0]=PP[1][0]=1;
	FOR(mm,100) {
		PP[0][mm+1]=PP[0][mm]*N;
		PP[1][mm+1]=PP[1][mm]*A;
		if(PP[0][mm+1]>A || PP[1][mm+1]>B || PP[0][mm+1]/N!=PP[0][mm] || PP[1][mm+1]/A!=PP[1][mm]) break;
	}
	cout << dpdp(A,B,mm) << endl;
	
}

まとめ

面白い問題だったけど、問題文が途中で変わったりいまいち。