PythonでAtCoder Beginner Contest 177 E - Coprime

AtCoder Beginner Contest 177に参加しました

D完。E問題で詰まったので復習としてまとめをします。

E - Coprime

本番中に考えたこと

実質「pairwise coprimeをどう判定するか？」という問題。
愚直にpairwiseで判定すると当然時間がかかるので素数ベースで考えるしかない。
要するに「すべての$A_i$について素因数が全くかぶらない」ということなので、$A_i$を素因数分解して出てきたものをカウント（p_cnt[$A_i$] に+1）していけばよさそう。
例：12→2・3、16→2、35→5・7なので、2が2回以上ありpairwiseではない
p_cntの最大値が1ならpairwise coprime、最大値がNならnot coprime、それ以外ならsetwise coprime。

ただし1を含む場合がコーナーケースなので、まず全部1ならpairwise coprime。
1, 1, 3, 6みたいに1以外がcoprimeでないならsetwise coprime、
1, 1, 3, 4みたいなケースならpairwise coprime。
したがって各$A_i$について考えるときに、それが1ならば個数をカウントしないようにする。

# Nの素因数分解を辞書で返す
def prime_fact(n):
    if n == 1:
        return {1: 1}

    root = int(math.sqrt(n))
    prime_dict = {}
    for i in range(2, root+1):
        cnt = 0
        while n % i == 0:
            cnt += 1
            n = n // i
        if cnt:
            prime_dict[i] = cnt
        if n == 1:
            break
    if n != 1:
        prime_dict[n] = 1
    return prime_dict


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

    if max(A) == 1:
        print("pairwise coprime")
        exit()

    p_cnt = defaultdict(int) # 素因数pが出てきた回数を記録
    for a in A:
        if a == 1: continue # 1の個数は数えてもしょうがないので無視
        AP = prime_fact(a)
        for apk in AP.keys():
            p_cnt[apk] += 1

    max_p = max(p_cnt.values())
    if max_p == 1:
        print("pairwise coprime")
    elif max_p == N:
        print("not coprime")
    else:
        print("setwise coprime")

1個（max_02.txt）だけTLEが取れない！！となりタイムアップ。

本番後に考えたこと

Cならともかく、Python(PyPy)では$N$が$10^6$もあると厳しい。
ここで、1,000,000以下の素数は78,498しかないことを利用することを思いつく。
すなわち$A_i$の種類($len(set(A))$)がだいたい8万以上の場合、鳩の巣原理からどこかの素因数に必ずカブリが発生してpairwise coprimeを満たさなくなる！！
なのでまず適当にgcdを取ってnot coprimeのケースを除外しておき、そのうえで8万種類以上あるケースをsetwise coprimeに突っ込むことで、Nを実質$10^6$から$10^5$くらいまで落とすことができる。

この枝刈りにより無事AC。

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

    if max(A) == 1:
        print("pairwise coprime")
        exit()

    # 全Aiの最大公約数が1より大きければnot coprime
    g = A[0]
    for a in A:
        g = math.gcd(g, a)
    if g > 1:
        print("not coprime")
        exit()

    # not coprimeでなく、Aが80000種類以上ならsetwise coprime
    SA = set(A)
    if len(SA) >= 80000:
        print("setwise coprime")
        exit()

    # 上記でないケースのみ頑張る
    p_cnt = defaultdict(int)
    cnt = 0
    for a in A:
        if a == 1: continue
        AP = prime_fact(a)
        for apk in AP.keys():
            p_cnt[apk] += 1

    max_p = max(p_cnt.values())
    if max_p == 1:
        print("pairwise coprime")
    else:
        print("setwise coprime")

より強いテストケースには負ける気がするのでおとなしく解説の高速化を使いましょう。