NoiminのNoise

競技プログラミング (多め) とWeb (たまに) ,自然言語処理 (ブログではまだ)。数式の書き方を一気に KaTeX に変えようとして記事を全削除してインポートし直すなどしたので,過去にブックマークされた記事は URL が変わってしまっている可能性があります…….

AtCoder Grand Contest 038 C - LCMs

問題: C - LCMs

公式解説: https://img.atcoder.jp/agc038/editorial.pdf

問題概要

長さ $ N (\leq 2 \times {10}^5) $ の数列 $ A_0, A_1, \cdots, A_{N-1} $ (ただし $ A_i \leq {10}^6 $ ) が与えられる.次式の値を 998244353 で割ったものを求めよ.

$ \sum_{i=0}^{N-2} \sum_{j=i+1}^{N-1} \mathit{lcm}(A_i, A_j) $

解法概要

参考: 公式解説pdf, 公式解説動画, tourist さんのソースコード

$ \mathit{lcm}(x, y) = \frac{xy}{\mathit{gcd}(x, y)} $

なので,

$ \sum_{d|\mathit{gcd}(x, y)} w_d = \frac{1}{\mathit{gcd}(x, y)} $ となるような係数 $ w_d $ が用意できると, $ A_0, A_1, \cdots, A_{N-1} $ の中で $ d $ の倍数となるものの和 (と二乗和) から求めたい値が求められるので嬉しい.

具体的には,以下のような式変形を行う:

$$ \begin{aligned} &\sum_{i=0}^{N-2} \sum_{j=i+1}^{N-1} \mathit{lcm}(A_i, A_j) \\ &= \sum_{i=0}^{N-2} \sum_{j=i+1}^{N-1} A_i A_j (\sum_{d|A_i, d|A_j} w_d) \\ &= \sum_{d=1}^{\max A_i} w_d \sum_{d|A_i, d|A_j, i \lt j} A_i A_j \\ &= \sum_{d=1}^{\max A_i} w_d \cdot \frac{1}{2} ((\sum_{d|A_i} A_i )^2 - \sum_{d|A_i} A_i^2) \end{aligned} $$

あとは式変形後の Σ の部分と,$ w_d $ が求められれば良いのだが,$ 1 + \frac{1}{2} + \frac{1}{3} + \cdots + \frac{1}{N} \simeq O(\log n)$ より時間計算量 $ O(n \log n) $ で求めることができる.

ただしこの問題は定数倍がきつめで,$ w_d $ を求めるときなどにオーダでは変わりなくても定数倍が重い処理を書いてしまうと簡単に TLE してしまうので注意.以下のソースコードでは TLE となる処理も AC だった処理と併記している.

ソースコード

#include <iostream>
#include <vector>

using namespace std;

using ll =  long long;

constexpr ll MOD = 998244353;
constexpr ll N_MAX = 200000;
constexpr ll A_MAX = 1000000;

inline ll modpow(ll a, ll t) {
    ll ret = 1LL;
    while(t){
        if(t & 1LL){
            ret *= a;
            ret %= MOD;
        }
        a *= a;
        a %= MOD;
        t >>= 1;
    }
    return ret;
}

inline ll modinv(ll a) {
    return modpow(a, MOD-2);
}

vector<ll> w(A_MAX);


// これだと TLE
// inline void init_w() {
//     w[1] = 1;
//     w[2] = (modinv(2)-1+MOD) % MOD;
//     for(int i=3;i<=A_MAX;++i) {
//         ll s = 0;
//         for(int j=1;j*j<=i;++j) {
//             if(i % j) continue;
//             s += w[j];
//             if(1 < j && j*j < i) s += w[i/j];
//             s %= MOD;
//         }
//         w[i] = (modinv(i) - s + MOD) % MOD;
//     }
// }
// こっちは余裕で AC
inline void init_w() {
    for(int i=1;i<=A_MAX;++i) {
        w[i] = modinv(i);
    }
    for(int i=1;i<=A_MAX;++i) {
        for(int j=i+i;j<=A_MAX;j+=i) {
            w[j] += MOD-w[i];
            w[j] %= MOD;
        }
    }
}

int main() {
    std::ios::sync_with_stdio(false); cin.tie(nullptr);
    init_w();
    int n;
    cin >> n;

    vector<ll> a(A_MAX+1, 0);
    ll aa;
    for(int i=0;i<n;++i) {
        cin >> aa;
        ++a[aa];
    }

    ll ans = 0LL;
    for(int i=1;i<=A_MAX;++i) {
        ll s = 0, t = 0;
        for(int j=i;j<=A_MAX;j+=i) {
            ll tmp = (j*a[j]) % MOD;
            s += tmp;
            s %= MOD;
            t += (tmp * j) % MOD;
            t %= MOD;
        }
        ans += (w[i] * (((s*s)%MOD - t + MOD) % MOD)) % MOD;
        ans %= MOD;
    }
    
    cout << (ans * modinv(2)) % MOD << endl;
}

所感

解説放送見て「頭いいー!」を連発することしかできなかった