TopCoder SRM603: PairsOfStrings

久しぶりに本番で通ったMedium。

A + C = C + B が成立するときは A は B を巡回シフトさせたものになる。
よって、もし A の文字列の最小周期が x だった場合は、対応するペアの個数は x になる。周期が x の文字列の個数は n % k == 0 のときは \(k^x\) 、それ以外の場合は \(0\)。

ただしこの値には約数の周期のパターン数も含まれているので、あとから取り除くようにする。

Practiceでは、BigInteger.ModPow を使って通してみた。

public class PairsOfStrings {
    const int mod = 1000000007;

    public List<int> Divisors(int x) {
        List<int> ret = new List<int>();
        for (int i = 1; i * i <= x; i++) {
            if (x % i == 0) {
                ret.Add(i);
                if (i != x / i) ret.Add(x / i);
            }
        }
        ret.Sort();
        return ret;
    }

    public int getNumber(int n, int k) {
        long ans = 0;
        var cyclePattern = new Dictionary<long, long>();
        var divisors = Divisors(n);
        foreach (int cycle in divisors) {
            long pattern = (long)BigInteger.ModPow(k, cycle, mod);
            foreach (int d in divisors) {
                if (cycle > d && cycle % d == 0) {
                    pattern = (pattern - cyclePattern[d] + mod) % mod;
                }
            }
            cyclePattern[cycle] = pattern;
            ans = (ans + pattern * cycle) % mod;
        }
        return (int)ans;
    }
}