けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 064 F - Rotated Palindromes (1000 点)

約数系包除。かなり教育的要素の多い問題だと思った!!!!!

  • 操作によって何通りできるのかを問う問題の取り組み方
  • 約数系包除の考え方

問題へのリンク

問題概要

 1 以上  K 以下の整数からなる数列のうち、以下のようにしてつくられるものは何通りあるか、1000000007 で割ったあまりで求めよ。

  • 回文数列になるようにする
  • 先頭要素を末尾いもっていくのを好きな回数だけ行う

制約

  •  1 \le N, K \le 10^{9}

考えたこと

わからないけど、こういうのは  K = 1 のとき、 K = 2 のとき...と順に感触を試していくうちに解ける気がする。

  •  K = 1 のとき、自明に  1 通りになる
  •  K = 2 のとき、 N が奇数だったら全パターン作れそう。 N が偶数だったら  1 2 がそれぞれ偶数個ずつのパターンは全部作れそう。
  •  K = 3 になると一気に難しくなる...

という感じになった。とりあえず  N の偶奇に意味がありそうな雰囲気が浮かび上がってる (後々実際に効いて来る)。でも、 K = 3 から見当もつかなくなったので、今回要求されている操作がどんなものか考えて見ることにした。

操作によって作られる個数

操作のやっていることは、数列を左右に分けて、左右を swap するというもの。そうすると考えたくなることは

  • 元の回文全体を、操作によって何種類の異なる数列が発生するかで分類する

ということをやりたくなる。この問題が単純に解けるならば、各  d に対して

  • 各回文が操作によって  d 種類の異なる数列を生じさせる数列が何通りあるかを数え上げて  f(d) 通りとする
  • そしてその操作によってできる数列は、通常重複があって、その重複度  k(d) が幾つになるかを考える (後でそれで割る)
  •  d について  f(d) × d/k(d) を合計する

をするだけで解けるはず。

f(d) を求める

まず、数列が操作を行うことで同じものになるなら、数列が繰り返し構造を持っていることを意味する (この考え方はメチャクチャ見る気がする)

具体的には、回文が  d 通りの数列を作るのなら、元の回文は周期  d にあっているはず。 \frac{N}{d} の長さの数列を  d 回繰り返す感じになっている。

よって  d N の約数に限られる。109 以下なら、最悪でも 100 個くらい。またそれぞれの繰り返し文字列もまた、回文でなければならないことも簡単に示せる。よって問題は


長さが  d の数列のうち、回文となっているものを数え上げよ。ただし繰り返し構造があってはならない。(その値が  f(d))


というものになる。いかにも包除原理になりそう。 d の各約数  e に対して、

  •  K^{\frac{e+1}{2}} ×  {\rm Mebius}(\frac{d}{e})

を合算する感じ。

重複度 k(d) を求める

 d が奇数だったら、操作によってつくられる  d 通りの数列はすべて異なる。よって  k(d) = 1

 d が偶数だったら、操作によってダブルカウントになることがわかる (123321 と 321123 など)。よって  k(d) = 2

まとめ

素直に考えれば素直に解ける気がする。

以上を実装して以下の通り

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

const int MOD = 1000000007;

vector<pair<long long, long long> > prime_factorize(long long n) {
    vector<pair<long long, long long> > res;
    for (long long p = 2; p * p <= n; ++p) {
        if (n % p != 0) continue;
        int num = 0;
        while (n % p == 0) { ++num; n /= p; }
        res.push_back(make_pair(p, num));
    }
    if (n != 1) res.push_back(make_pair(n, 1));
    return res;
}

vector<long long> calc_divisor(long long n) {
    vector<long long> res;
    for (long long i = 1LL; i*i <= n; ++i) {
        if (n % i == 0) {
            res.push_back(i);
            long long j = n / i;
            if (j != i) res.push_back(j);
        }
    }
    sort(res.begin(), res.end());
    return res;
}

long long modpow(long long a, long long n, long long mod) {
    long long res = 1;
    while (n > 0) {
        if (n & 1) res = res * a % mod;
        a = a * a % mod;
        n >>= 1;
    }
    return res;
}

int main() {
    long long N, K;
    cin >> N >> K;
    
    vector<long long> div = calc_divisor(N);
    long long res = 0;
    for (auto d : div) {
        long long tmp = 0;
        vector<long long> div2 = calc_divisor(d);
        for (auto e : div2) {
            auto pf = prime_factorize(d/e);
            bool ok = true;
            for (auto p : pf) {
                if (p.second > 1) ok = false;
            }
            if (!ok) continue;
            int mebius = (pf.size() & 1 ? -1 : 1);
            long long pow = modpow(K, (e+1)/2, MOD);
            tmp += pow * mebius;
            tmp = (tmp % MOD + MOD) % MOD;
        }
        res += tmp * (d % 2 == 0 ? d/2 : d) % MOD;
        res = res % MOD;
    }
    cout << res << endl;
}