いむらや競プロ雑記

キーワード:競プロ、育児、石川県

ABC249 D - Index Trio

atcoder.jp

カウント対象が成立するためには、左辺が割り切れる必要がある(Aのそれぞれが整数のため)ことから、iについて全探索するとした場合に、Aiの約数を列挙しておくとよさそう。さらに、AjとAk相当は、まじめに探索しなくても、個数から引いてこれそうな気持ちになる。

約数列挙は、元の数が大きい場合に大変になりそうな気がするが、同じ数に対して計算をしなくていいように、とりあえずdefaultdictで持ってみることにした。

あとは、Aiを決めたときにあり得る、仮想的なAjとAkを全探索し、そのパターンが実際いくつあるかを数える。
Ajになり得るのはAiの約数であり、更にAiとAjを決めたときにAkになり得るのは、Ai//Ajである。i,j,kの重複はし放題のため、何ら条件は付けずに、これらの個数をかけ合わせればOK 。(Aに存在しないAjやAkは、個数が0で返ってくるので、変にカウントしてしまうこともない)

※なお、約数列挙はこちらから拝借しました qiita.com

from collections import defaultdict, deque


def make_divisors(n):
    lower_divisors, upper_divisors = [], []
    i = 1
    while i * i <= n:
        if n % i == 0:
            lower_divisors.append(i)
            if i != n // i:
                upper_divisors.append(n // i)
        i += 1
    return lower_divisors + upper_divisors[::-1]


def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    dy = defaultdict(list) #yakusuu
    dk = defaultdict(int) #kosuu

    for a in A:
        dk[a] += 1
        if dy[a] == []:
            temp = make_divisors(a)
            dy[a] = temp

    ans = 0
    for _i, v in dy.items():
        inum = dk[_i]
        for _j in v:
            _k = _i // _j
            jnum = dk[_j]
            knum = dk[_k]

            ans += knum * jnum * inum

    print(ans)


resolve()