TopCoder SRM 603 Div1 Medium PairsOfStrings

問題

n文字でアルファベットの最初のk種類からなる文字列A, Bのうち、
次の条件を満たす(A, B)のペアの総数をmod 10^9 + 7で求めよ。

ある文字列Cが存在し、A + C = C + Bとなる。
（ただし + は文字列の結合を表す）
AまたはBが一文字でも違うとき、二つのペアは異なると数えるものとする。

制約条件

n≦10億
k≦26

方針

A + C = C + Bになるということは、
Bの先頭何文字かを切り取ってBの末尾につけると、Aに一致するということ。
文字列Aの末尾を切り取ってBを何通り生成できるかを、重複なく数えられればよい。

一つのAに対して、違う位置で切り取って、同じBが生成されてしまう必要十分条件を考えると、
文字列Aの周期がTであるときに、n/T回同じ切断で同じ文字列が生成されることがわかる。

周期がTの文字列をそれぞれ何通りあるか包除原理で数え、
切断箇所T通り倍するのを、それぞれのTについて足し合わせればいい。

ソースコード

const int mod = (int)1e9 + 7;
ll pw(ll n, ll m){
	ll res = 1;
	for(; m; m /= 2){
		if(m & 1) res = res * n % mod;
		n = n * n % mod;
	}
	return res;
}
map<int, int> dp;
int calc(int n, int k){
	if(dp.count(n)) return dp[n];
	
	ll res = (pw(k, n) + mod - dp[1]) % mod;
	for(int i = 2; i * i <= n; i++) if(n % i == 0){
		
		res = (res + mod - calc(i, k)) % mod;
		if(i * i != n) res = (res + mod - calc(n / i, k)) % mod;
	}
	return dp[n] = res;
}

class PairsOfStrings {
	public:
	int getNumber(int n, int k) {
		
		dp.clear();
		dp[1] = k;
		calc(n, k);
		ll ans = 0;
		each(i, dp){
			ans += (ll)i->first * i->second % mod;
		}
		return ans % mod;
	}
};