エラトステネスの篩をPythonで実装して関数化する

概要

以前この記事で
hirohirohirohiros.hatenablog.com
エラトステネスの篩は実装しました．しかし，以前のコンテストでこの関数では対応出来ない問題に当たりました．
　このエラトステネスのふるい関数はiが素数の時リストi番目がTrue，素数でないときFalseが入ったリストを返す関数です．しかし，単に素数のみが入ったリストが欲しい場合もあります．このような場合に対応するために，関数を書き換えます．
　get_primesというbool型の変数を新たに引数に用意します．なお，既に使われているコードに新たに書き換える必要を無くすためにget_primes=FalseとデフォルトではFalseとなるようにしています．これがTrueの時，新しい処理を追記します．
　今まで，あらかじめ全てが素数だと仮定しておき，素数出なかったらFalseにするという処理にしていたので，while文の条件式はp^2

pythonコード

def erato_sieve(x, get_primes = False):
    """エラトステネスの篩

    0からxまでの素数を返す
    計算量: O(n loglog n)

    Args:
        x (int): 素数判定をしたい数の最大値
        get_primes(bool): 素数のみが入ったリストを返すかどうか

    Returns:
        list: get_primesがTrueの時
                素数飲みが入ったリスト
              get_primesがFalseの時
                i番目が素数ならTrue, そうでないならFalse
    """
    primes = []
    prime_list = [True]*(x+1)
    prime_list[0] = False
    prime_list[1] = False
    p = 2
    while p <= x:
        if not prime_list[p]:
            p += 1
            continue
        
        primes.append(p)
        i = p*2
        while i <= x:
            prime_list[i] = False
            i += p
        p += 1

    return primes if get_primes else prime_list