AtCoder ABC152 E の復習

AtCoder ABC152 E 問題の復習です。この問題、方針としては LCM を \( A_i \) の最小公倍数とすると、

\( A_i B_i = LCM \)

となるのは自明なはずです。これは \( B_1 + B_2 + … B_N \) を最小化したいなら \( A_i B_i \) を最小化できればいいためです。

通常 LCM は A / B * (gcd(A, B)) (gcd は A, B の最大公約数を計算する関数) で求まりますが、今回の問題では制約上 LCM が 64 bit に収まらない可能性があります (多倍長整数を使えば計算できますし、実際 python 等はこのままいけたっぽいですが)。

そこで LCM そのものを求めるのではなく、素因数とそのカウントのような、つまり素因数分解した形で LCM を表現することにします。これは 3 を 3 ^ 1、4 を 2 ^ 2、12 であれば (2 ^ 2) * (3 ^ 1) と見るような考え方です。データ構造としては例えば map が使えます。

あとは実際に計算していくだけ… なんですけど、自分が時間内に気付かなかった、知らなかったテクニックが解説動画に一杯あったのでそれらをいつかの自分のためにまとめます。主に参考にしたのは以下の 2 つです。

ちなみに最終的な自分の解答はこちらです。何だかどんどんライブラリが増えていく…

素因数分解

ある正整数 n を素因数分解する場合、単純なやり方としては

\( \sqrt n \) 以下の素数を列挙
それらで n を割ってみる

すればよく、C++ で書くと以下のようになります。

const int MAX_N = 1000000000;
vector<bool> numbers;
vector<int> primes;
vector<int> factors;

// エラトステネスの篩で sqrt(n) 以下の素数を列挙する
void collect_primes(const int n) {
    numbers.resize(n, true);
    for (int i = 2; i * i <= n; i++) {
        if (numbers[i]) {
            primes.push_back(i);
            for (int j = i; j * j <= n; j += i) {
                numbers[j] = false;
            }
        }
    }
}

// n について素因数分解
void factorize(const int n) {
    int x = n;
    for (int i = 0; i < (int) primes.size(); i++) {
        int prime = primes[i];
        while (x % prime == 0) {
            factors.push_back(prime);
            x /= prime;
        }
    }
    // A prime factor larger than sqrt(x) is added here.
    if (x != 1) {
        factors.push_back(x);
    }
}

// AOJ の素因数分解を行う問題
// http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=NTL_1_A
int main(void) {
    int n;
    cin >> n;

    collect_primes(n);

    factorize(n);
    cout << n << ": " << factors[0];
    for (int i = 1; i < (int) factors.size(); i++) {
        cout << " " << factors[i];
    }
    cout << endl;

    return 0;
}

時間計算量はステップ 1 (collect_primes) で \( O(\sqrt n) \)、ステップ 2 (factorize) は正確な評価はできないのですが最悪の場合素数の個数分 (primes のサイズ) の時間がかかると思います。

この問題の場合各 \( A_i \) について素因数分解する必要があり、ステップ 1 については一度やれば使い回せますが、ステップ 2 については N 回行う必要があります。

このように多数の正整数について素因数分解しなければならない場合、AtCoder Beginner Contest 152 - AtCoder Live や AtCoder ABC 152 E - Flatten (500 点) - けんちょんの競プロ精進記録で紹介されている方法が使えます。詳細はリンク先に譲りますが、エラトステネスの篩を利用して素数を求めつつ各整数の素因数分解も \( O(素因数の個数) \) で行うことができます。今回の問題制約では素数が 80000 ぐらい存在するので、愚直に素因数分解を行うよりも多少早くなるはずです。実装例は AtCoder Live のライブラリだったり冒頭にリンクした自分の提出コード内 EratosthenesSieve クラス内にあります。

余談ですがエラトステネスの篩による表の作成にかかる時間計算量は \( O(N \log \log N) \) らしいです。これ自体の説明はできないのですが、\( O(N \log N) \) の気持ちというのはどうやら

\( \sum_{i=1}^{N} \frac{N}{i} = N \sum_{i=1}^{N} \frac{1}{i} \) となり、シグマの部分が

\( \sum_{i=1}^{N} \frac{1}{i} \approx \int_{1}^{N} \frac{1}{x} dx = [\log x]_{1}^{N} = \log N \) と近似できるということみたいです。

説明としては調和級数1+1/2+1/3…が発散することの証明の積分を用いた証明が視覚的にわかりやすいと思いますが、グラフで示される通り近似はできるけど抑え込めるわけではないはずです。

(追記)

\( \log N + 1 \) で抑え込めます。

mod 演算上での除算

この問題で求めるべきは

\( \frac{LCM}{A_1} + \frac{LCM}{A_2} + … + \frac{LCM}{A_N} \bmod 1000000007 \)

です。

LCM mod 1000000007 は既に LCM を素因数分解した形がわかっているので加算、乗算のみで求められます。 mod は加算、乗算については計算途中で mod を取りつつ計算してよい、というのは競プロの比較的簡単な問題でもよく見ます。

問題は除算ですが、a mod p について

p が素数
a と p が互いに素

の場合には、フェルマーの小定理により mod p 上での a の逆数が求まる、つまり a による除算が可能です。フェルマーの小定理から

\( a^{p-1} \equiv 1 \bmod p \)

であり、これは

\( a \times a^{p-2} \equiv 1 \bmod p \)

なので mod p 上での a の逆数が \( a^{p-2} \) になるということです。

このあたりの mod を使用した計算については「1000000007 で割ったあまり」の求め方を総特集！〜逆元から離散対数まで〜が参考になるかと思います。

またこれで除算を乗算のみで計算できるわけですが、p は 1000000007 と大きい数なので愚直に計算すると恐らく TLE になります。大きな冪乗を計算するテクニックとして例えば蟻本では P.114 で繰り返し二乗法と呼ばれる方法が紹介されており、これを使用することで \( O( \log N ) \) に時間計算量が抑えられます。

typedef long long ll;

// assume that n >= 0;
ll mod_pow(const ll x, ll n, const ll mod) {
    ll res = 1;
    if (n == 0) return 1;
    if (n & 1) res = x;
    ll y = mod_pow(x, n >> 1, mod);
    res = (res * (y * y % mod)) % mod;
    return res;
}

// AOJ の冪数を求める問題
// ref. http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=NTL_1_B
int main(void) {
    const ll mod = 1000000007LL;
    ll x, n;
    cin >> x >> n;

    cout << mod_pow(x, n, mod) << endl;

    return 0;
}