【ABC152】E - Flatten

atcoder.jp

解法

$ A_i B_i $はなるべく小さい方が数列Bの総和は小さくなる。
数列Aの要素にそれぞれ適当な値をかけて、全て同じ値を作り出したいのだから、$ A_i B_i $は数列Aの要素の最小公倍数にすると最小になることがわかる。

最小公倍数はとんでもなく大きな数になる可能性がある。
C++には最小公倍数を求める関数が標準でそわっているが、modを取らないとオーバーフローしてしまう。
そこで、独自に最小公倍数をとることにする。
具体的には数列の要素をそれぞれ素因数分解して、各素因数毎に個数の最大値をとって、最後にそれを掛け合わせる。

最大公倍数をA[i]で割ったものの総和が答えになる。

実装

osa_k法を使うと高速に素因数分解できる。
下の記事を参考にした。

blog.manuel1024.com

vector<int> sieve(int n){
    vector<int> min_factor(n+1);

    for(int i = 0; i <= n; i++) min_factor[i] = i;

    for(int i = 2; i*i <= n; i++){
        if(min_factor[i] == i){
            for(int j = 2; i*j <= n; j++){
                if(min_factor[i*j] > i){
                    min_factor[i*j] = i;
                }
            }
        }
    }
    return min_factor;
}

vector<int> factor(vector<int> &min_factor, int m){
    vector<int> ans;
    while(m > 1){
        ans.push_back(min_factor[m]);
        m /= min_factor[m];
    }
    return ans;
}

void solve(long long N, std::vector<long long> A){
    auto min_factor = sieve(1000010); // 事前準備

    map<int, int> primes; // 素因数の最大数を格納
    REP (i, N) {
        // 素因数分解 O(logA[i])
        auto p = factor(min_factor, A[i]);
        map<int, int> mp; // A[i]の素因数を格納
        for (auto m : p) chmax(primes[m], ++mp[m]);
    }

    // 最小公倍数を計算
    mint LCM = 1;
    for (auto p : primes) LCM *= modpow(mint(p.first), p.second);

    // 答えを計算
    mint ans = 0;
    REP (i, N) ans += LCM / mint(A[i]);
    cout << ans << endl;
}