htkb-proconの日記

初心者がPythonで問題解いた記録

yukicoder No.407 鴨等素数間隔列の数え上げ

No.407 鴨等素数間隔列の数え上げ

素数の列挙を行う必要のある問題。素数列挙といえばエラトステネスのふるいで、これを普通に組んでTLEになった問題に今まで当たったことはなかったけどこの問題ではだめだった。

定数倍に気を遣ったエラトステネスのふるいでの素数列挙
https://yukicoder.me/submissions/260831

上のコードで気を遣ったポイントは

  • 最初から偶数を省いて素数候補を作る
  • ふるいで素数の倍数を落とす際に、素数の偶数倍は偶数で↑により既に候補に含まれないため、i*3からi*2飛ばしでrangeを作る(今見ているのが3ならば9, 15, 21, 27...みたいな数列を作ってふるい落とす)
  • set.difference_updateは差集合を得る、つまり+=と同じだがset以外のiterableも受け付けるのでsetインスタンスを作り捨てずに済む
  • もちろん属性アクセスはループ前にキャッシュ

でそこそこ速いつもりだったんだけど普通に通らなかったので高速な素数列挙を調べてみたところ、とりあえずエラトステネスの前処理でお手軽にできるWheel factorizationとやらが見つかったので試してみた。とても分かりやすかったサイト↓

理屈は上のサイト見れば全部わかるのであとは実装どうすっかだけど、pythonでforループでちまちま要素追加してくと余計遅くなりそうなので、とりあえずサイクル長=6で素数候補は{2, 3, 5から6飛ばしの数、7から6飛ばしの数}と決め打ちでsetを作る。ふるい落としのターンでは5から6飛ばしの数と7から6飛ばしの数のみチェックすればいいので、for i in range(5, N+1, 6)みたいにして回して、ii+2に対してチェックしていけばいい。上のサイトの長方形の数列を見て理屈を理解すればライブラリ化してなくてもその場で問題なく素で書ける。

でACしたコード https://yukicoder.me/submissions/260853

微妙な改善って感じだけど……。サイクル長を30にしたら実装が煩雑になって手元では余計に遅くなった。そもそもset素数候補全部持つのが重いんだろうか。もっと根本的に軽い実装あるかなぁ