[AtCoder][競技プログラミング][python] ABC 249 D-Index Trio

AtCoder

Photo by Nothing Ahead on Pexels.com

2022.05.09

Table of Contents

問題

D - Index Trio

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

方針

\(1\)個ずつ確認すると\(O(N^3)\)かかる。素因数を列挙する方針にした。素因数は以下のコードで列挙できる。計算量は\(O(\log{N})\)。

import math
def prime_factors(n):
    factors = []
    m = int(math.sqrt(n))+1
    for i in range(1, m):
        if not n % i:
            factors.append(i)
            if i != n//i:
                factors.append(n//i)
    return factors

後はCounterで要素を管理し、足し上げていく。Counterは要素がない場合、きちんと\(0\)を出力してくれるので便利。

解答

#atcoder template
def main():
    import sys
    input = sys.stdin.readline
    #文字列入力の時は上記はerrorとなる。
    #ここにコード
    #input
    n = int(input())
    a = list(map(int, input().split()))

    #output
    import math
    def prime_factors(n):
        factors = []
        m = int(math.sqrt(n))+1
        for i in range(1, m):
            if not n % i:
                factors.append(i)
                if i != n//i:
                    factors.append(n//i)
        return factors

    from collections import Counter
    b = Counter(a)
    c = list(b.keys())
    answer = 0

    for i in c:
        temp = prime_factors(i)
        for j in temp:
            answer += b[i]*b[j]*b[i/j]

    print(answer)

    #N = 1のときなどcorner caseを確認！
if __name__ == "__main__":
    main()