D - Index Trio

[hxrxchang]

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

問題概要

数列Aがある。 A[i] / A[j] == A[k] となる (i, j, k) の組み合わせを求めよ。 i, j, kにはAの範囲内であること以外に制限はない。 (すべて同じでも、単調増加列になっていなくてもいい。)

解き方

Aの各要素A[i]の約数を求めておく。 Aの各要素の出現回数も求めておく。

入力例1 で考える。

3
6 2 3

i を6とすると、6の約数は 1, 2, 3, 6 であり、これらの各要素を j とする。 するとk は i // j なので 6, 3, 2, 1 になる。 i, j, k がAの中に何回出現するかそれぞれ掛け合わせると、i = 6 であるときの組合せ数が求められる。 i を Aの各要素にして探索して足し合わせたのが答え。

計算量は

• 約数列挙がroot(N) なので、全体で N * root(N)
• 約数をそれぞれ見ていくのも 10 ** 5 以下の数字でもっとも約数の多い数は 83160の128個 (ChatGPT調べ) なので

from collections import defaultdict

def calc_divisors(N):
    # 答えを表す集合
    res = []

    # 各整数 i が N の約数かどうかを調べる
    for i in range(1, N + 1):
        # √N で打ち切り
        if i * i > N:
            break

        # i が N の約数でない場合はスキップ
        if N % i != 0:
            continue

        # i は約数である
        res.append(i)

        # N ÷ i も約数である (重複に注意)
        if N // i != i:
            res.append(N // i)

    # 約数を小さい順に並び替えて出力
    res.sort()
    return res

N = int(input())
A = list(map(int, input().split()))

dividors = defaultdict(list)
cnt = defaultdict(int)

for a in A:
    cnt[a] += 1
    if dividors[a] == []:
        dividors[a] = calc_divisors(a)

ans = 0
for i in dividors.keys():
    for j in dividors[i]:
        k = i // j
        ans += cnt[i] * cnt[j] * cnt[k]

print(ans)

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

[hxrxchang]

Ai / Aj = Ak は Ai = Ak * Ak と式変形したほうが約数列挙の問題だということがわかりやすい。