はまやんはまやんはまやん

hamayanhamayan's blog

Index Trio [モノグサプログラミングコンテスト2022(AtCoder Beginner Contest 249) D]

https://atcoder.jp/contests/abc249/tasks/abc249_d

前提知識

解説

https://atcoder.jp/contests/abc249/submissions/31211388

とある数xの約数列挙をO(sqrt(x))で行う方法がある。
もし分からない場合は「約数列挙」あたりで検索して、勉強してくる必要がある。
(今回は、naiveに約数列挙するより、区間で約数列挙を事前計算でしておいた方が早そうな気もするが…未検証)

何から考えるか

まずi,j,kの全列挙を考えてみる。これはO(N3)で当然間に合わない。
だが、経験則から、こういった添え字の列挙問題はどれか1つは全列挙して、他を最適化する方針が多い。
とりあえずiを全列挙して、他をうまく効率化できないか、使える性質を探してみよう。

約数

条件で、A[i] / A[j]というのがあるが、計算結果がA[k]になるかを考えるので、
この分数は計算すると整数になる必要がある。
iを全探索、つまり、A[i]を固定して他を最適に考えるとすると、A[j]というのはA[i]の約数である必要が最低限ある。
しかも、A[j]が定まればA[k]も一意に定まるので、実質A[i]が固定化されているときの選択肢はA[i]の約数個数分しかない。
これは使えそうな性質である。

約数を列挙する方法については、高速なやり方があるので、それを利用する。
約数の個数は問題にならないかという疑問があるが、「高度合成数」という概念があり、
2*105だと160くらいが上限っぽいのでこれなら問題ない。

解法へ

以下を前計算しておく。
cnt[x] := A[i]=xとなるiの個数

iを全列挙してA[i]を固定化させたら、次は条件を満たすjを列挙する。
これは約数だけでいいので、代わりに、jの全列挙はA[i]の約数の全列挙をする。

とあるA[i]の約数divとなるjの個数はcnt[div]個であり、A[k]はA[i]/divとなるのでkの個数はcnt[A[i]/div]個となる。
jとkの組はどんな組でも問題ないので、cnt[div] * cnt[A[i]/div]通りが
その約数を使った時の条件を満たす組み合わせとなる。

あとは全部総和を取れば答え。

int N, A[201010];
int cnt[201010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> A[i];
    rep(i, 0, N) cnt[A[i]]++;

    ll ans = 0;
    rep(i, 0, N) {
        auto divs = enumdiv(A[i]);
        fore(div, divs) {
            int cntJ = cnt[div];
            int cntK = cnt[A[i] / div];
            ans += 1LL * cntJ * cntK;
        }
    }
    cout << ans << endl;
}