素数判定いろいろ - AKS素数判定法

数学好きのエンジニアしてます @srtk86 です。本日の日付 1223 は素数です。

素数かどうかを判定するアルゴリズムについて色々見てみたので、その話をしようと思います。

[前回](https://qiita.com/srtk86/items/609737d50c9ef5f5dc59)は素数の性質の1つであるフェルマーの小定理と、そこから導き出される確率的な素数判定アルゴリズムについて書きました。
今回は、決定的な判定法の1つである「AKS素数判定法」について見ていきたいと思います。

## 目次

[シンプルな素数判定と、素数の分布](https://qiita.com/srtk86/items/874639e361917e5016d4)
[フェルマーテスト・ミラーラビン素数判定法](https://qiita.com/srtk86/items/609737d50c9ef5f5dc59)
AKS素数判定法 **- 今ここ！**

## PRIMES is in P！

AKS素数判定法は、前回紹介したフェルマーテストなどとは異なり、**100%の確率** で素数を判定できるアルゴリズムの1つです。**決定的アルゴリズム** などとも呼ばれます。登場は2002年ですが、2005年までにいくつかの修正を経て、[こちらの論文(PDF)](https://www.cse.iitk.ac.in/users/manindra/algebra/primality_v6.pdf) にまとめられています。
論文のタイトルが、何よりこのアルゴリズムの特徴を示しています。その名も **PRIMES is in P**!!

ここでの「P」とは、[「P vs. NP 問題」](https://ja.wikipedia.org/wiki/P%E2%89%A0NP%E4%BA%88%E6%83%B3) における「（クラス）P」を指しています。クラスPは、「**多項式時間で解ける**問題」の集合とも言われます。
つまりAKS素数判定法は **素数判定の問題がクラスPに属する、つまり多項式時間で解ける** ことを示した最初のアルゴリズムなのです。かっこいいですよね...！

1点、誤解しやすい点として、「多項式時間で判定できる」というときの多項式時間は、「**入力サイズ** $m$」 に関する多項式で表せる時間 ですので、注意が必要です。
[最初の記事](https://qiita.com/srtk86/items/874639e361917e5016d4) で紹介した、$2$から$n-1$ までで割っていくだけのアルゴリズムも、「入力される値 $n$」 に対してはおよそ $n - 2$ 回程度の計算回数（1次式）に収まりますが、「入力サイズ $m$」 に対してで考えると、（入力サイズをビット数とした場合は） だいたい $n = 2^{m}$ なので、計算回数も $2^{m}$ 回程度の指数関数時間となり、多項式の枠には収まりません。
逆に「入力サイズ$m$」に対する多項式時間は、「入力される値 $n$」 に対してでいうと、$\log{n}$ の多項式で表せる必要があるので、ハードルはかなり高いです。

## AKS素数判定法のアルゴリズム

アルゴリズムは以下の通りです。$n$は2以上の自然数とします。

1. $n = a^{b} \quad (a, b \in \mathbb{N}, b > 1) \quad$ という形で表せるなら、$n$ は**合成数**
2. $o_r(n) > \log^2{n}$ となるような最小の数$r$を求める
3. $a \leq r$ なる $a$ のいずれかが $1 < (a, n) < n$ を満たす場合、$n$ は**合成数** （ここでの $(a, n)$ は、$a$ と $n$ の最大公約数）
4. $n \leq r$ ならば、$n$ は**素数**
5. $1 \leq a \leq [\sqrt{\phi(r)} \log{n}] $ の範囲の$a$に対して
    - $(X + a)^{n} \not\equiv X^{n} + a \pmod{X^{r} - 1, n}$ ならば、$n$ は**合成数**
6. $n$ は**素数**

2.にある $o_r(n)$ は**位数**と呼ばれるもので、$n^{e} \equiv 1 \pmod r$ を満たす最小の自然数$e$を指します。
また、5.にある $\phi(r)$ は**トーシェント関数**あるいは**オイラー数**と呼ばれるもので、$1$から$r$までの中で、$r$との最大公約数が$1$になるものの個数を指します。

上で紹介した論文の記述の大半は、このアルゴリズムで素数が100%判定できることの証明と、このプロセスの計算量が $\log^{10.5} {n}$ 以下のオーダーに収まることの証明に費やされています。

ではPythonプログラムに変換していきます。難しい箇所が幾つかあるので、ステップごとに見ていきます。

### 1. 累乗数の判定

> $n = a^{b} \quad (a, b \in \mathbb{N}, b > 1) \quad$ という形で表せるなら、$n$ は**合成数**

累乗数の形に限定すれば、素因数分解は少ない回数で試せます。$b$ になりうる値は $\log_2{n}$ 以下になるので、$2$から$\log_2{n}$ の範囲の $b$ に対して $n$ の $b$ 乗根を計算し、ちょうど整数になるものがあれば $n$ は累乗数、すなわち合成数と判定することができます。

プログラムにするとこんな感じになります。（誤差の調整で2回判定してしまっています...）

```python3
# nが累乗数かどうかを判定
def is_perfect_power(n):
    size = int(math.log2(n) + 1)
    for b in range(2, size):
        a = int(math.pow(n, 1 / b))
        if a ** b == n or (a+1) ** b == n: return True
    return False
```


### 2. 十分大きな位数を持つ数

> $o_r(n) > \log^2{n}$ となるような最小の数$r$を求める

$n^{e} \pmod r$ を計算していき、$1$になったら $e$ と $\log^2{n}$ を比較。$e$ の方が大きければ、そのときの $r$ が求めるものです。
条件がやや複雑ですが、一つ一つの計算はわりと単純です。

このような $r$ は、高々 $\log^5{n}$ のオーダーの値までには見つかる、というのが論文でも証明されており、$r = n - 1$ まで試行が走ることは殆どありません。（$n$が小さい場合には例外あり）

```python3
# 十分大きな位数を持つ最小の数
def enough_order_modulo(n):
    a = int(math.log(n) ** 2)
    for r in range(1, n):
        order = 0
        prod = 1
        for e in range(1, r):
            prod = prod * n % r
            if prod == 1:
                order = e
                break
        if order > a:
            return r
    return n
```

### 3. 純粋な素数判定

> $a \leq r$ なる $a$ のいずれかが $1 < (a, n) < n$ を満たす場合、$n$ は**合成数** （ここでの $(a, n)$ は、$a$ と $n$ の最大公約数）

シンプルな素数判定を、わる数を $r$ までに制限して実施すれば良いだけです。

### 4. 前のステップで判定が完了する場合

> $n \leq r$ ならば、$n$ は**素数**

単純に比較するだけです。

### 5. トーシェント関数および剰余多項式の計算

> - $1 \leq a \leq [\sqrt{\phi(r)} \log{n}] $ の範囲の $a$ に対して
>     - $(X + a)^{n} \not\equiv X^{n} + a \pmod{X^{r} - 1, n}$ ならば、$n$ は**合成数**

$\phi(r)$ については、以下の公式を使用して計算します。

```math
\phi(r) = r \prod_{p | r} \left( 1 - \frac{1}{p} \right) 
```

コードにするとこんな感じ。$r$ を割り切る素因数を洗い出し、その全てについて $\displaystyle 1 - \frac{1}{p}$ を $r$ にかけていく流れです。

```python3
# トーシェント関数
def tortient(r):
    ps = primes(r)
    res = r
    for p in ps:
        res  = res * (p - 1) / p
    return res

# rの素因数のリスト
def primes(r):
    n = r
    res = {}
    for p in range(2, int(math.sqrt(r)) + 1):
        while n % p == 0:
            res.add(p)
            n /= p
    return res
```

続いて剰余多項式の計算ですが、配列を用いて計算してみます。

次数が$i$の部分の係数を、配列のインデックス$i$の位置に入れるようにします。たとえば2次式 $x^2 + 2$ は、`[2, 0, 1]` という配列に変換されます。(0次の係数は2, 1次の係数は0, 2次の係数は1なので)

次は配列で多項式の積を表現する必要があります。
$f(x)$の$i$次の項と、$g(x)$の$j$次の項の積は $i+j$ 次になります。なので原則として、各項の積は $i+j$ 番目のインデックスに入っていくことになります。
そこに $X^r - 1$ の剰余が加わると、$r$次より上の項は、そのまま$r$個分ずれる感じになるので、各項の積は $i+j \pmod r$ 番目のインデックスに入ります。

あとは配列の累乗ですが、積の表現さえできていれば、繰り返し自乗法の考え方にしたがって自然に再帰的に書くことができます。

上記の思考をプログラムに落とすと、こんな感じになるでしょうか。(係数の mod $ n$ の計算も、特に説明してはいませんが、プログラムには含めています)

```python3
class PolynomialModulo(object):

    def pow(self, ls, n, r):
        self.ls = ls
        self.n = n
        self.r = r

        return self.__pow(self.n)

    def __pow(self, m):
        if m == 1: return self.ls
        if m % 2 == 0:
            pls = self.__pow(m // 2)
            return self.__product(pls, pls)
        else:
            return self.__product(self.__pow(m - 1), self.__pow(1))

    def __product(self, ls1, ls2):
        res = [0] * min(len(ls1) + len(ls2) - 1, self.r)
        for i in range(len(ls1)):
            for j in range(len(ls2)):
                res[(i + j) % self.r] += ls1[i] * ls2[j]
        l = len(res)
        for k in reversed(range(l)):
            res[k] %= self.n
            if k == len(res) - 1 and res[k] == 0: res.pop(k)
        return res
```

これを利用して、実施したい判定を行う関数をつくります。

```python3
# (x + a)^n と x^n + a が mod x^r - 1, n で合同かどうかを判定
def is_congruent(a, n, r):
    p = PolynomialModulo()
    ls1 = p.pow([a, 1], n, r)

    i = n % r
    ls2 = [0] * (i + 1)
    ls2[0] = a % n
    ls2[i] = 1

    return ls1 == ls2
```

`ls1` が `(x + a)^n`, `ls2` が `x^n + a` にあたります。

### 6. nは素数

> $n$ は**素数**

$n$は素数です。

### アルゴリズムの実装の全体

ここまで見てきたプログラムを繋げた、素数判定プログラムの全体像がこちらです。

```python3:aks_prime_test.py
import math
import numpy as np

def is_prime(n):
    if n == 1: return False

    # Step 1
    if is_perfect_power(n): return False

    # Step 2
    r = enough_order_modulo(n)

    # Step 3
    for a in range(2, min(r+1, n)):
        if n % a == 0:
            return False

    # Step 4
    if n <= r: return True

    # Step 5
    for a in range(1, int(math.sqrt(tortient(r)) * math.log(n)) + 1):
        if not is_congruent(a, n, r):
            return False

    # Step 6
    return True

# nが累乗数かどうかを判定
def is_perfect_power(n):
    size = int(math.log2(n) + 1)
    for k in range(2, size):
        pn = int(math.pow(n, 1 / k))
        if pn ** k == n or (pn+1) ** k == n: return True
    return False

# 十分大きな位数を持つ最小の数
def enough_order_modulo(n):
    a = int(math.log(n) ** 2)
    for r in range(1, n):
        order = 0
        prod = 1
        for e in range(1, r):
            prod = prod * n % r
            if prod == 1:
                order = e
                break
        if order > a:
            return r
    return n

# オイラーのトーシェント関数
def tortient(r):
    ps = primes(r)
    res = r
    for p in ps:
        res  = res * (p - 1) / p
    return res

# rの素因数のリスト
def primes(r):
    n = r
    res = set()
    for p in range(2, int(math.sqrt(r)) + 1):
        while n % p == 0:
            res.add(p)
            n /= p
    return res

# (x + a)^n と x^n + a が mod x^r - 1, n で合同かどうかを判定
def is_congruent(a, n, r):
    p = PolynomialModulo()
    ls1 = p.pow([a, 1], n, r)

    i = n % r
    ls2 = [0] * (i + 1)
    ls2[0] = a % n
    ls2[i] = 1

    return ls1 == ls2

class PolynomialModulo(object):

    def pow(self, ls, n, r):
        self.ls = ls
        self.n = n
        self.r = r

        return self.__pow(self.n)

    def __pow(self, m):
        if m == 1: return self.ls
        if m % 2 == 0:
            pls = self.__pow(m // 2)
            return self.__product(pls, pls)
        else:
            return self.__product(self.__pow(m - 1), self.__pow(1))

    def __product(self, ls1, ls2):
        res = [0] * min(len(ls1) + len(ls2) - 1, self.r)
        for i in range(len(ls1)):
            for j in range(len(ls2)):
                res[(i + j) % self.r] += ls1[i] * ls2[j]
        l = len(res)
        for k in reversed(range(l)):
            res[k] %= self.n
            if k == len(res) - 1 and res[k] == 0: res.pop(k)
        return res

if __name__ == '__main__':
    for n in range(1, 101):
        print(str(n) + ': ' + str(is_prime(n)))

```

ただ、注意が1点。これ、**ものすごく遅いです(笑)**
10桁いかないぐらいの素数判定にも数分単位の時間を取るため、実行時は十分注意してください。（実行を途中で止めたいときは ctrl + C です）

## 課題

多項式の計算のところについては、高速Fourier変換を使うともう少しパフォーマンスが出せるらしいのですが、今回はそこまではできませんでした。
適切な型を構築していって円分体の計算を構成して、その上で走らせてみるとかは、実現できると面白そうなので、気力が向かえばやってみたいなぁと思います。

## まとめ

- AKSはかっこいい
- AKSはおそい
- でもAKSはえらい