SRM603 Div1 Medium - PairsOfStrings

TopCoder SRM

包除原理🤤

問題

以下の条件を満たす文字列のペア(a,b)の個数を求めよ
1) aとbの長さはどっちもn
2) aもbもアルファベットの最初のk文字のみからなる
3) a+c = c+bとなるような文字列cが存在する

1 <= n <= 10^9
1 <= k <= 26

解法

3つ目の条件は、
(aの先頭x文字) = (bの末尾x文字) かつ (aの末尾n-x文字) = (bの先頭n-x文字)
となるようなx (0 <= x <= n)が存在すること と言い換えられる
(これを満たすときはc=(aの先頭x文字)とすれば成り立つ)
つまりbはaを0〜n-1回rotateしたものであればいい
aはk^n通り考えられ、bは1個のaに対してn通り考えられる
ただaをn未満回rotateした時点でaに戻るようなときはbを重複してカウントしてしまうから、その分をいい感じに引いてやる必要がある

aをd(＜n)回rotateすると元に戻るってのは、例えばabcabc(d=3)とかbb(d=1)みたいに、aに周期があるようなときである
こういうaはk^d通り存在し、それぞれに対応する重複しないbはd個(aを0〜d-1回rotateしたもの)存在する
このdはnの約数全通りあり得る
これらより答えは包除原理の要領で
Σ [k^d - Σ(k^s)] * d (d: nの正の約数全て、 s: dの正の約数のうちd以外のもの)
となる
例えば周期6のものは一部周期1・2・3のものと重複してるからその分を引く、みたいな感じ

10^9以下の正整数の約数の個数は最大で1344個しかないから計算量はその2乗でも大丈夫
modpowで更にlognが付くけどまあ大丈夫

class PairsOfStrings {
public:
	int getNumber(int n, int k) {
		vector<int> d;
		for (int i=1; i*i<=n; i++) if (n%i==0) {
			d.push_back(i);
			if (i*i!=n) d.push_back(n/i);
		}
		sort(all(d));
		ll m = d.size(), res = 0, num[2020] = {};
		rep(i,m) {
			num[i] = mop(k,d[i]);
			rep(j,i) if (d[i]%d[j]==0) (num[i] += mod - num[j]) %= mod;
			(res += num[i] * d[i]) %= mod;
		}
		return res;
	}
};

https://community.topcoder.com/stat?c=problem_statement&pm=12788&rd=15836