30歳で競プロに目覚めた霊長類のブログ

チンパンジーと一般人のあいだ

ABC177-E coprime 4つの解法

ABC177-E coprimePython or PyPyで間に合う解法を4つ記します

AtCoder Problems では緑difficultyらしいのですが(!?)コンテスト中には解けずに苦汁を飲みました

問題概要

  • N 個の整数が与えられる
  • N 個の整数の異なるすべてのペアが互いに素 ( \gcd(A_i, A_j) = 1 ) のとき、pairwise coprime と出力する
  • そうでなくてすべての整数の \gcd1 のとき、setwise coprime と出力する
  • いずれでもないとき、not coprime と出力する
  •  2 \le N \le 10^{6}, 1 \le A_i \le 10^{6}

コンテスト中の思考

  • not coprime は最初に全体の \gcd を取ればok。 (2 以上なら not coprime)
    • (追記) 以下紹介する4つの解法は最初に↑の判定を済ませている前提で書きます
  • 本質は pairwisesetwise の判定
  • A の異なる要素に共通の素因数が見つかったら setwise 確定。
    • とりあえず全部素因数分解するのが思いつく
    •  O(N \cdot \sqrt{A_i}) ぐらいかかってしまうと思い敬遠。1
      • →実は間に合う(解法3)
  • 全部素因数分解、を避ける方法で判定しなきゃいけない! (←思い込み)
  • 前計算で 10^{6} までの素数を求めておいてうまくゴリゴリやろう...! (←沼の始まり)

解法1 高速素因数分解

writer の kyopro_friends さんの想定解です。
エラトステネスの篩の進化版みたいなやり方で、osa_k法 とも呼ばれているみたいです。2

MAXN = 10**6+10
sieve = [i for i in range(MAXN+1)]
p = 2
while p*p <= MAXN:
    if sieve[p] == p:
        for q in range(2*p,MAXN+1,p):
            if sieve[q] == q:
                sieve[q] = p
    p += 1
  • まず配列を用意し
    • [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15, ...] のようにします 3
  • 次に2の倍数で書き換えられていないところを2で書き換えます
    • [0,1,2,3,2,5,2,7,2,9,2,11,2,13,2,15,...]
  • 次に3の倍数で書き換えられていないところを3で書き換えます
    • [0,1,2,3,2,5,2,7,2,3,2,11,2,13,2,3,...]`
  • 次は4ですが すでに2に書き換えられているのでスキップします
  • 次に5の倍数で書き換えられていないところを5で書き換えます
  • ...

のようにして配列を作っておくと、素因数分解がめっちゃ高速にできます!!

tmp = set()
while a > 1:
    tmp.add(sieve[a])
    a //= sieve[a]

素因数の列挙は a > 1 である限り sieve[a] で割り続けるだけ!
(今回は列挙だけなので set でやっていますが、もちろん dictCounter を使うことで素因数の個数も数えられます)

例えば a=60 として、

  • sieve[60] == 2 なので 2 を抽出し、a2 で割る。
  • sieve[30] == 2 なので 2 を抽出し、a2 で割る。
  • sieve[15] == 3 なので 3 を抽出し、a3 で割る。
  • sieve[5] == 5 なので 5 を抽出し、a5 で割る。
  • a==1 になったのでおわり。 tmp = {2, 3, 5} が抽出された。

計算量は O(A \log \log A + N \log A) のようです
配列の前計算が O(A \log \log A)素因数分解O(N \log A) です
素因数分解めっちゃ速いですね!!

ACコードの全体はこちら Python: 834ms

解法2 等間隔に舐めて調和級数で間に合うやつ

これは解説放送ですぬけさんが説明されていた方法です
詳しくは放送をご覧いただきたいのですが、なんと 素因数分解素数判定も必要ありません

MAXA = 10**6
B = [0] * (MAXA+1)
for a in A:
    if a==1: continue
    if B[a]:
        print('setwise coprime')
        exit()
    B[a] = 1

まずヒストグラムみたいのを作ります (上記の B です)
同じ数値が2度出現したら setwise 確定で exit() しています。
ただし 1 が複数個あっても \gcd(1,1) = 1 なので setwise とは確定しないことに注意! (コーナーケース!)
結果として B の要素としては 01 のいずれかが入る形になります。

for n in range(2,MAXA+1):
    c = 0
    for m in range(n,MAXA+1,n):
        c += B[m]
        if c > 1:
            print('setwise coprime')
            exit()

さてここが肝です。上のコードでは

  • 2,3,4,5, ... 1000000 までのそれぞれを n として、
    • n から始めて 2n, 3n, ... (これを mとする) の B[m] をチェックしています
    • n のどれかで、B[m] が 2回以上カウントされたら setwise 確定です。

一見  10^{6} までのループを2重に回しているので間に合わなそうに見えるのですが、実は間に合います
よく見ると1回の外側のループでは MAXA // n 回ぐらいしか処理されません。
 O(\displaystyle \sum_{i=1}^{A} \frac{A}{i}) は調和級数という考え方で  O(A \log A) になりますよというのが知られています (詳しくは解説放送をご覧ください)

ACコードの全体はこちら(PyPy) 327ms と、こちら(Python) 1320ms※ 4

解法3 愚直に素因数分解しちゃう

えー!間に合っちゃうんですか!

ライナスさんにTwitterでリプライ頂きました
なんと愚直に試し割りして素因数分解しても間に合っちゃうみたいです!

計算量を雑に見積もると  O(N \sqrt{ \max A})10^{9} オーダーなんですが、ちゃんと考えると1桁以上落ちるようです
というのも、

  • 10^{6} 以下の素数78498 個で、 10^{6} に比べて1桁以上少ない
  • 鳩ノ巣原理という考え方を適用すると、(追記)2以上のN78499 個以上存在すると自動的にどれかしらの素因数が重複する! setwise 確定 !!
  • 最悪ケースを一応考えると
    • 999983 のような素数がならぶ → 1000 までの試し割りでok
    • 2n乗、 3m乗、 5l乗、、、みたいに (l,m,n\max A を超えないような最大の指数) が並ぶことも一応あるけどこれは早く検出できるうえに指数はせいぜい20以下

となり、78499 \cdot 1000 (追記) + 10^{6}程度で  8 \cdot 10^{7} に収まるオーダーです
これでPyPyで間に合いました

私も書いてみました(Pythonは間に合いませんでした)
ACコードの全体はこちら PyPy: 678ms

解法4 ゴリゴリスペシャ

きれいな解法ではないのでおまけです
コンテスト中に 2WA までこぎつけたのですがコーナーケースでやられていました
コンテスト後にねちょさんにコーナーケースを見つけていただきました!

先にコンテスト後のACコードを貼ります (下記2つは同一コードです)
PyPy 425ms、Python 1986ms (!!)

やったこと

  • 全体の \gcd を取って 2 以上なら not coprime
  • 重複チェック1 以外で 2 つ以上重複している要素があれば setwise coprime
    • 1 以外で の部分は本番中で見つけられなかったコーナーケースです...
  • エラトステネスの篩で 10^{6} ぐらいまでの素数を検出する。実際に使ったのはそのうちの 10^{3} 以下の素数
  • A_i について、
    • 10^{3} 以下の素数 で割れる限り割っていく
      • その過程で割れた素数は記録する
      • 重複していたら setwise coprime
    • 1 以外になったら 10^{3}より大きい素数なので記録する
      • 重複していたら setwise coprime
  • いずれのチェックにもかからなければ pairwise coprime

という方法を取りました
ちょっとでも試し割りの回数を減らそうと素数を列挙したのですが、結局解法3で述べた鳩ノ巣原理 による影響の方が大きくて結局計算量的には間に合っています (本番中は鳩ノ巣原理も見破っていなかったです)


おわりです!!
1問でたくさん学びがある良い問題でした!!!!


  1.  \sqrt{A_i} \le 10^{3} ぐらいまで試し割りする方法しかないと思い込んでしまった。。PyPyは基本的にPythonより速いが  10^{8} オーダーは間に合うけど  10^{9} オーダーは間に合わない肌感。

  2. おさけ?

  3. 0,1は素数じゃないので扱い注意です (-1などで埋めるのが良い?)

  4. リンク先を見ると分かると思うんですが、PyPyと同一のコードではTLEしてしまいました。簡易的なPythonの高速化として「全体を関数で囲む」というのがあるんですが、それをしたらACしました!