【Python】n以下の素数を求めるプログラムをPythonで作ってみた！

概要

人工知能(AI)の研究者であり，著者には「人工知能は人間を超えるか」などがある東京大学の松尾豊先生の松尾研究室が開講している「グローバル消費インテリジェンス寄附講座」の演習パートのコンテンツがJupyterNotebbok形式で公開されている．
今回は，そのChapter1の総合問題に関連する問題を解いてみた．

動作環境

• Windows10(64bit)
• Python 3.7.2

問題

問題は以下のようなものである．(※文面はそのままではない)

n以下の素数を全て求めるプログラムを作れ．

それほど難しいものではなく，基礎的な良問ですので，一度自分で考えてみてください．

解答

素数とは，1より大きい自然数で1とその数自身のみを約数にもつ数のことである．または簡単に約数が2個の自然数と言うこともできる．
ゆえに，2以上自分以下の数で自分を割っていき，全て割り切れなければ素数ということになる．これをそのままプログラムにしてもよいが，これだとかなり無駄があり遅いプログラムになってしまう．工夫する方法は様々あるが，今回は素数判定の有名なアルゴリズムであるエラトステネスのふるいという方法で実行してみる．

コード

# nまでの素数を表示させるプログラム
import math 

def sieve_of_eratosthenes(n): 
    candidate = [i for i in range(2, n+1)] 
    prime = []
    limit = math.sqrt(n) + 1 

    while True:
        p = min(candidate) 
        if limit <= p:
            prime.extend(candidate) 
            break
        prime.append(p) 

        candidate = [i for i in candidate if i % p != 0]

    print(prime)

n = int(input('n=?\n'))

sieve_of_eratosthenes(n)

解説

ではコードの解説をしていきます．
まずは，最初の部分．

    candidate = [i for i in range(2, n+1)] 
    prime = []
    limit = math.sqrt(n) + 1 

素数の候補をcandidateリストに入れます．ここではリスト内包表記を使いました．1は素数ではないので最初から除外します．また，range(start, end)はstratからend-1までということに注意をしてください．
2行目で素数の空リストを作っておき，素数と判定できたらこのリストに追加していきます．
素数は半分だけ考えればよいことが数学的に示されているので$\sqrt{n}$まで調べます．ここで+1をつけないと平方数を判定するときに平方数を素数と判定してしまうので+1を付けておきます．ceil関数やfloor関数を使って+1をしたりすればlimitはintにできますが，今回はこのままでやることにします．

続いて，While文の前半部分です．

        p = min(candidate) 
        if limit <= p:
            prime.extend(candidate) 
            break
        prime.append(p) 

pにcandidateの最小値を指定します．一番最初なら2ですね．次のループでは3,5,7…となっていきます．
pがlimit以上であれば，そのときに残っているのは全て素数なのでprime.extend()でprimeリストとcandidateリストを結合してwhile分を抜け出します．
そうでなければ，prime.append(p)でcandidateの最小値をprimeリストに登録します.
もちろんこれは一番最初なら2．その次のループでは3,5,7…となっていきます．

最後にWhile文の後半部分です．

        candidate = [i for i in candidate if i % p != 0]

ここは少しわかりにくいかもしれません．自分で一度手を動かしてみると理解しやすいです．
candidateのなかでpの倍数であるものを削除しています．一番最初なら2の倍数は素数でないので削除しています．
この部分は以下のように記述しても構いません．

        i = 0
        while True:
            if i >= len(candidate): 
                break
            elif candidate[i] % p == 0: 
                candidate.pop(i) 
            i += 1      

やっていることは同じです．candidate.pop(i)はつまり，pで割ったあまりがゼロであれば．candiate[i]は素数ではないのでcandidateリストからpop(削除)します．

コメントを元に改善したコード

ありがたいことに，Qiitaユーザーの方から有益なフィードバックをいただいたのでそれをもとにコードを改善してみました！

# nまでの素数を表示させるプログラム

def sieve_of_eratosthenes(n): #エラトステネスのふるい
    candidate = list(range(2, n+1)) #候補
    prime = []
    limit = n**0.5 + 1 #+1がないと平方数を素数にしてしまう

    while True:
        p = candidate[0] 
        if limit <= p:
            prime.extend(candidate) #この時点での残っているのは素数
            break
        prime.append(p) #pを素数リストに登録

        candidate = [i for i in candidate if i % p != 0]

        # i = 0
        # while True:
            # if i >= len(candidate): #iはcandidateの長さ分
                # break
            # elif candidate[i] % p == 0: #余りがゼロなら
                # candidate.pop(i) #素数ではないのでpop(取り除く)する
            # i += 1

    return prime

n = int(input('n=?\n'))

prime = sieve_of_eratosthenes(n)

print(prime)

主な変更点は4カ所です．
1.リスト内包表記ではなく，リスト関数に変更した
2.math.sqrt()ではなく**0.5にすることでmathライブラリを使わずに済むようにした
3.candidateリストは小さい順に並んでいるので，p=candidate[0]とした．
4.sieve_of_eratosthenes()関数を素数データを返すようにした．

以上の変更点について詳しく見ていきます．

1.
変更前

candidate = [i for i in range(2, n+1)]

変更後

candidate = list(range(2, n+1))

リスト内包表記を使わなくてもlist()を使えば十分でした．この辺の計算速度とかを比べようとすると，各関数がどのようなアルゴリズムで動いているかとかデータ構造とかをきちんと知る必要がありそうですね．

2.
変更前

limit = math.sqrt(n) + 1

変更後

limit = n**0.5 + 1

mathライブラリをインポートしなくて済むようになるのでこれは，こちらの方が速そうですね．またこれに伴い最初のmathライブラリをインポートする所のコードも削除しています．
言ってみたら，大学生になって三次方程式の解の公式(カルダノの公式)を知った．けど公式は長すぎて覚えてはいない．三次方程式を解きたい！となったとき，大学の教科書を開いてカルダノの公式があるところを「どこだどこだ？」と探すよりも，因数を見つけて因数分解して，二次方程式に帰着させた方が，たぶんラクで速いし，教科書を一冊分持っていかなくて済みますよね？(もちろんカルダノの公式自体はとても素晴らしいものです．)

3.
変更前

p = min(candidate)

変更後

p = candidate[0] 

candidateリストは昇順に並んでいるので，min()を使う必要はありませんでした．
言ってみたら，学校で出席番号が一番の子の解答を，番号順に並べられた解答用紙の束からすべて確認して，五十音で一番なのはどの子だと探しているようなもんですね．
とても無駄ですね．

4.
変更前

print(prime)

変更後

return prime

オリジナルでは関数の返り値を持たせずに，primeリストを返していました．primeを表示させることだけが目的ならこれでも問題はないといえばないのですが，
関数はprimeリストを返すようにして，データを受け取った側でそのデータをprintしたり，個数を数えたりと，他の処理にも使えるようにした方が汎用性の高い関数になりますね．
言ってみたら，カメラ，ゲーム機，電卓，携帯電話を別個で持ち歩くより，iPhoneを一台持ち歩いた方が汎用性が高くて便利ですよね．(もちろん，たとえばカメラにはiPhoneにはない良さがたくさんありますが，今は機能はほぼ同じと仮定しています．)

それに伴い，以下の部分も追加しました．

prime = sieve_of_eratosthenes(n)
print(prime)

primeはリストデータなので，そのままprintしてあげればよいです．型変換は必要ありません．文字列連結するときなどは型変換をする必要があります．

まとめ

今回は，非常に基礎的な素数を求めるプログラムを作ってみました．素数を求めるアルゴリズムはたくさんあり，計算速度の議論も興味深いので，そちらについてもまた触れたいと思います．
[追記]
コードに関する建設的なフィードバックももらえると大変助かります！