TopCoder SRM 603 Div1 Medium PairsOfStrings

EasyをやらかしたおかげでMediumを終えることができなかった。
わかってしまえば簡単な問題。
http://community.topcoder.com/stat?c=problem_statement&pm=12788

問題

先頭K種類以内のアルファベットだけを使って作られるN文字の文字列A,Bを考える。
N,Kに対し、ある文字列Cを用いてA+C=C+Bとなるような(A,B)の対の数を答えよ。

解法

まずA+C=C+Bとなる条件を考える。
CをL文字とすると、Aの先頭L文字がBの末尾L文字と一致し、Aの後半(N-L)文字がBの先頭(N-L)文字と一致する。
すなわちBはAを回転させた文字列となる。

AはN文字なのでK^N通り作ることができ、Bはそれを0～(N-1)文字のN通り回転させて生成すると合わせてN*(K^N)通りの対が作れる。

ただし、Aがもともと同じパターンの繰り返しの場合、0～(N-1)文字回転させる間にBとして複数回同じ文字列が生成されるため、その分は除かないといけない。

Aの最短周期がNの約数であるD文字であるような場合を考える。
AをD文字をN/D回繰り返して生成する方法はK^D通りあるが、周期がDの約数Eであるようなケースはその型のぞかないといけないのでである。また、Aから生成できるユニークなBは0～(D-1)のD通り回転させたものであるので、先の値にDを掛ければよい。

これまでの検討をまとめると、以下の手順となる。

• Nの約数をd(N)個列挙する。O(√N)で可能。
• 上記Nの約数Dを最短周期とするN文字の文字列のパターンをF(D)とすると、Dの小さい順にで求めることができる。O(d(N)^2)で処理可能。
• が答え。

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

vector<ll> enumdiv(ll n) {
	set<ll> S;
	for(ll i=1;i*i<=n;i++) if(n%i==0) S.insert(i),S.insert(n/i); //include 1
	S.insert(n);
	return vector<ll>(S.begin(),S.end());
}

ll numnum[2000];

class PairsOfStrings {
	public:
	
	int getNumber(int n, int k) {
		ll ret=0;
		int x,y;
		
		vector<ll> S=enumdiv(n);
		FOR(x,S.size()) {
			numnum[x]=modpow(k,S[x]);
			FOR(y,x) if(S[x]%S[y]==0) numnum[x]+=mo-numnum[y];
			numnum[x]%=mo;
			ret = (ret+numnum[x]*S[x])%mo;
		}
		
		return ret%mo;
	}

};

まとめ

落ち着いてやればコードもさほど難しくない。
Easyで完全に混乱していた…。
ついでに約数列挙ルーチンもライブラリ化しておいた。