【高速】素数列挙・素因数分解・約数総数・約数列挙のアルゴリズム【Python】

AtCoderなどの整数論の問題でよく使うアルゴリズムを, 関数として実装します. コンテスト中でもすぐに参照・ピペできるよう, 簡潔かつ高速で動く実装を心掛けています.

1. 整数N以下の素数列挙

エラトステネスの篩を使って, 0～nまでの整数に対する素数判定を行う.
n + 1個のTrue or Falseで表された要素を出力する.
その中から, Trueのものだけを取り出してリストに格納すれば, 素数列挙できる.

import math

# エラトステネスの篩
def prime(n):
    is_prime = [True] * (n + 1)
    is_prime[0], is_prime[1] = False, False
    for i in range(2, int(math.sqrt(n)) + 1):
        if is_prime[i]:
            for j in range(2 * i, n + 1, i):
                is_prime[j] = False
    return is_prime

# 素数列挙
N = int(input())
# M以下の整数の素数判定
prime = prime(N)

# 素数判定を基に, 素数リストを作成
prime_list = []
for i in range(N + 1):
    if prime[i]:
        prime_list.append(i)
        
print(*prime_list)

例として, 100以下の素数をすべて列挙してみる.

in: 100
out: 2 3 5 7 11 13 17 19 23 29 31 37 41 43 47 53 59 61 67 71 73 79 83 89 97

2. 素因数分解

2から\(\sqrt{N}\)までの数で割り切れるか判定する.
割り切れる場合, 割り切れなくなるまでその数で割り続け, 割った回数をCounterオブジェクトに記録していく.
小さい数から割っていくので, 必ず素数で割ることになる.
割る数として, エラトステネスの篩（ふるい）で先に素数を列挙しておくのもあり.
\(O(\sqrt{N})\の計算量で済む.

import math
from collections import Counter

# 素因数分解を行う関数
def prime_factorization(n):
    # Counterクラスを呼び出してオブジェクト生成
    counter = Counter()
    
    # 素数で割り切れるかの判定
    for p in range(2, int(N ** 0.5) + 1):
        # １回割り切れるごとに+1カウントしていく
        while n % p == 0:
            counter[p] += 1
            n //= p
    # n が1より大きい数字として残っていれば、素数
    if n != 1:
        counter[n] += 1
    return counter

N = int(input())
result = prime_factorization(N)
for key, value in result.items():
    for i in range(value):
        print(key, end = " ")

例として, 250を素因数分解してみる.

in : 250
out : 2 5 5 5

3. 約数の個数

ある数Nを次のように素因数分解できるとき、Nの約数の総数は第２式で表せる．
\begin{align*}
& N = x_1^{q_1}x_2^{q_2}\cdots x_n^{q_n}\\
& \displaystyle \prod_{i=1}^n (q_i + 1) = (q_1 + 1) \times (q_2 + 1) \times \cdots \times (q_n + 1)
\end{align*}

素因数の個数の管理には, Counter()を使用する.
Counter()オブジェクトは, defaultdict同様, 生成していないキーを指定しても, 参照結果として値0が返されるので便利.
2. と同様, エラトステネスの篩（ふるい）で先に素数を列挙しておくのもあり.

import math
from collections import Counter

# 素因数分解を行う関数
def prime_factorization(n):
    # Counterクラスを呼び出してオブジェクト生成
    counter = Counter()
    
    # 素数で割り切れるかの判定
    for p in range(2, int(N ** 0.5) + 1):
        # １回割り切れるごとに+1カウントしていく
        while n % p == 0:
            counter[p] += 1
            n //= p
    # n が1より大きい数字として残っていれば、素数
    if n != 1:
        counter[n] += 1
    return counter

N = int(input())
result = prime_factorization(N)
tot = 1
for key, value in result.items():
    tot *= value + 1
print(tot)

4. 約数の列挙

引数nに対して，約数をリストで返す関数を実装する．

\(1 \le i \le \sqrt(n)\)の範囲で\(n / i\)が割り切れるかを探索．
割り切れた場合，(\i\)を約数としてリストに格納していく．

def divisor(n):
    ds = []
    for i in range(1, int(n ** 0.5) + 1):
        # nをiで割り切れる場合
        if n % i == 0:
            # iは約数
            ds.append(i)
            # n = i^2ではない場合
            if i != n // i:
                ds.append(n//i)
    ds.sort()
    return ds

実行例