　この記事はwathematicaの企画である「Wathematica Advent Chalender2025」の12月13日分です.
　今回はAKSアルゴリズムという素数判定法について書きました.実装に関しては誰がみても理解できるように丁寧に書いたつもりなので是非見てください.

1.はじめに

整数を見た時にその数が素数かどうかを考えてしまうのは人間の性だと思います.しかし,ある大きな数が素数かどうかを判定するのはとても難しい問題だと思います.そこで活躍するのが素数判定法であり,現在では様々な素数判定法が発表されています。今回は,その一つであるAKSアルゴリズムの実装と計算が決定多項式時間で収まることの証明をしたいと思います.

2.AKSアルゴリズムとは

2002年に史上初の“決定的多項式時間”素数判定法であるAKSアルゴリズムが発表されました.発表論文のタイトルが「PRIMES is in P」であるように,この発表のすごいところは素数判定がクラスPに属すると示された点にあると思います.ここでクラスPとは簡単に言うと,判定問題のなかで多項式時間以内に解けるアルゴリズムが存在するものの全体のことです.

アルゴリズムの流れ

まずはこのアルゴリズムの流れを紹介します。ここで判定したい整数を$n$とする.

$n$が整数の冪乗か調べる.(冪乗なら合成数のため終了)
$Ord_{r}(n)>{(log_2 n)}^{2}$を満たす最小の$r$を求める.
$1\leqq a\leqq r$に対して$a$と$n$が互いに素か調べる(そうでないなら合成数のため終了)
$n\leqq r$なら素数.そうでなければ次に行く
$a=1$から$a=\lfloor \sqrt{\phi (r)}log_2 n\rfloor$までの$a$について${(x+a)}^{n}\equiv x^{n}+a\pmod{x^{r}-1,n}$が成り立てば素数

過程で出てきた関数の定義を確認します.

$Ord_{r}(n)$･･･ $n^{m}\equiv 1\pmod r$となる最小の正の整数$m$
$\phi (n)$･･･$n$と互いに素な$n$以下の正の整数の個数(Eulerの$\phi$関数)

さらにSTEP5の$\pmod{x^{r}-1,n}$の考え方がわかりにくいと思うので解説すると,各項の係数は常に$\pmod{n}$で考えることを前提として$x$の指数$s$が$r$以上の項の係数は$x^{s\pmod{r}}$の項の係数に足してあげるということです.
この流れで素数判定を行うことができます.次にアルゴリズムを実装してみましょう.

3.AKSアルゴリズムの実装

まずは先ほど確認した関数を実装します.
その後先ほどのせた1-5の手順に則って実装していきます.

利用する関数

Ord関数

def multiplicative_order(n,r):
    if gcd(n,r)!=1:
        return 0
    k=1,x=n%r
    while x!=1:
        x=(x*n)%r
        k+=1
        if k>r:
            return 0
    return k

最初のif分岐は$n$と$r$が互いに素でない場合は$order$が定義できないためです。(何回$r$をかけても$1$にならない)
そしてwhileで$n$に$n$をかけ続け$r$で割った余りが$1$になった時点で終了します.$k$がループの回数を記録しているのでその値が戻り値です.

Eulerのφ関数

def phi(n):
    count = 0
    for i in range (1, n):
        if gcd(n, i)==1:
            count + = 1
    return count

今回mathモジュールのgcdを用いました.
これらをもとにAKSアルゴリズムを実装していきます.

STEP1 冪乗か判定

まずは入力された整数nがある自然数の冪乗かどうかを次に記す関数で判定していきます.

def is_perfect_power(n):
    if n <= 1:#1以下の数の除去
        return False
    max_b = int(math.log2(n))#bとして考えられる最大値
    for b in range(2, max_b + 1):
        a = round(n ** (1 / b))
        if a**b or (a+1)**b ==n:
            return True
    return False

指数bを$2〜log_{2} n$まで変えながら$a=n^{\frac{1}{b}}$を求め$a^{b} = n$となる$a$が存在すればTrue,全部の$b$で失敗したらFalseという流れです.
入力$n$が整数の冪乗ならTrue,そうでないならFalseを返す関数にしました.
つまりFalseなら次のステップに行くと考えてください.
浮動小数誤差でズレることを考慮してif a**b or(a+1)**b==n:と書いています.

STEP2 条件を満たす最小のrを求める

$Ord_{r}(n)>{(log_2 n)}^{2}$となる最小の$r$を求める関数は次のとおりです.

    limit=int(math.log2(n)**2)#不等式の右側
    max_r=int(math.log2(n)**5)+2#取り得る最大のr+2

    r=3#取り得る最小のr
    while r <= max_r:
        if gcd(n, r) != 1:
            continue
        if multiplicative_order(n,r)>limit:
            break
        r+=1
    else:
        return False

max_rとして$r$の上限を設定していますが,これは条件を満たす$r$が$max[3,({log_2 n})^5]$までに存在すると「PRIMES is in P」内で証明されているからです.この証明はこの後紹介します.

STEP3 aとnが互いに素か

    for a in range(2, r + 1):
        g = gcd(a, n)
        if 1<g<n:
            return False

$a=1$の時,$gcd(a, n)=1$であることは明らかなのでループを$a=2$から始めています.

STEP4 rがn以上なら素数

    if n <= r:
        return True

STEP5

ここがこのアルゴリズムのメインとなる部分です.
まず最後の判定をするために次の関数を導入します.

def poly_mul_mod(p, q, mod, r):
    res = [0] * r
    for i, a in enumerate(p):
        for j, b in enumerate(q):
            res[(i + j) % r] = (res[(i + j) % r] + a * b) % mod
    return res

poly_mul_mod 関数は、多項式$p(x)$と$q(x)$の積を法$x^r - 1$と法 mod のもとで計算する関数です。
2つの多項式の係数を走査し、次数が r を超える項は$(i+j) \bmod r$ によって折り返します（つまり $x^r \equiv 1$ として扱う）。

def poly_pow_mod(base, exp, mod, r):
    res = [1] + [0] * (r - 1)
    x = base[:]
    while exp:
        if exp & 1:
            res = poly_mul_mod(res, x, mod, r)
        x = poly_mul_mod(x, x, mod, r)
        exp >>= 1
    return res

poly_pow_mod 関数は、多項式の累乗を法 $x^r - 1$ と法 mod のもとで高速に計算する関数です.
二分累乗法（繰り返し二乗法）を用いることで指数が大きい場合でも計算量を大きく抑えられます.
これらの関数を用いて行う最後のチェックは次です.

    max_a = int(math.sqrt(phi(r)) * math.log2(n))
    for a in range(1, max_a + 1):
        base = [a % n, 1] + [0] * (r - 2)
        left = poly_pow_mod(base, n, n, r)

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

        if left != right:
            return False

    return True

この部分では
$(x + a)^n \equiv x^n + a \pmod{(x^r - 1,\ n)}$
が条件を満たす $a (1 \le a \le \lfloor \sqrt{\varphi(r)}\log_2 n \rfloor)$について成り立つかどうかを確かめています。
baseは多項式 $x + a$ を、right は $x^n + a$ を$(\mathbb{Z}/n\mathbb{Z})[x]/(x^r - 1)$ 上の係数配列として表現したものです。
poly_pow_mod(base, n, n, r) により $(x + a)^n$ を同じ環の中で計算し,どれか一つでも等しくなければ合成数と判定し（False を返す),すべての $a$ で一致すれば素数と判定して True を返すという仕組みです。

これで素数判定が終了しました.最後にコード全体を紹介します.

コード全体

import math
from math import gcd

#Ord関数
def multiplicative_order(n,r):
    if gcd(n,r)!=1:
        return 0
    k = 1
    x = n % r
    while x != 1:
        x = ( x * n ) % r
        k += 1
        if k > r:
            return 0
    return k
#オイラーのトーシェント関数
def phi(n):
    count=0
    for i in range (1,n):
        if gcd(n,i)==1:
            count+=1
    return count

def is_perfect_power(n):
    if n <= 1:
        return False
    max_b = int(math.log2(n))
    for b in range(2, max_b + 1):
        a = round(n ** (1 / b))
        if a ** b == n:
            return True
    return False

def poly_mul_mod(p, q, mod, r):
    res = [0] * r
    for i, a in enumerate(p):
        for j, b in enumerate(q):
            res[(i + j) % r] = (res[(i + j) % r] + a * b) % mod
    return res

def poly_pow_mod(base, exp, mod, r):
    res = [1] + [0] * (r - 1)
    x = base[:]
    while exp:
        if exp & 1:
            res = poly_mul_mod(res, x, mod, r)
        x = poly_mul_mod(x, x, mod, r)
        exp >>= 1
    return res

def is_prime(n):
    if is_perfect_power(n):
        return False

    limit=int(math.log2(n)**2)
    max_r=int(math.log2(n)**5)+2

    r=2
    while r <= max_r:
        if gcd(n, r) != 1:
            r+=1
            continue

        if multiplicative_order(n, r)>limit:
            break

        r+=1
    else:
        return False

    for a in range(2, r + 1):
        g = gcd(a, n)
        if 1<g<n:
            return False

    if n <= r:
        return True

    max_a = int(math.sqrt(phi(r)) * math.log2(n))
    for a in range(1, max_a + 1):
        base = [a % n, 1] + [0] * (r - 2)
        left = poly_pow_mod(base, n, n, r)

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

        if left != right:
            return False

    return True

nums = [2, 3, 5, 7, 11, 21, 97, 121, 15934531]
for x in nums:
    print(x, "is ", "prime" if is_prime(x) else "composite")

これまでコードの解説を行ってきましたが,あまりこの素数判定は早くないんじゃないかと思った方はいると思います.実はその通りで手持ちのノートPCでは8-9桁の素数判定が限界です.しかしこのアルゴリズムが偉いのはクラスPに属していることです.その証明が気になっている方も多いと思うので論文での流れに沿って証明を追ってみましょう.

3.AKSアルゴリズムで素数が判定できることの証明

まず最初に次の二つの補題を示しておきます.
STEP2において最小の$r$を求めましたが上限があることを前提としました.この二つの補題はその上限に関するものです.

Lemma 1

$m\geq7$が整数なら,$1,…,m$の最小公倍数は$2^{m}$以上である

(proof)
ここで,ベータ関数とガンマ関数を定義します.$a>0,b>0$において

$B(a,b)\stackrel{\mathrm{def}}{=}\int_0^1x^{a-1}(1-x)^{b-1}dx$
$\Gamma(x)\stackrel{\mathrm{def}}{=}\int_0^\infty t^{x-1}e^{-t}dt$

また,ベータ関数をガンマ関数を用いて

$B(a,b)=\frac{\Gamma(a)\Gamma(b)}{\Gamma(a+b)}$

と表すことができます.
この関数を用いて証明していきます.
$1…m$の最小公倍数を$d_{m}$とおきます.
この$d_{m}$をベータ関数を用いて評価していきます.
$1\leq a\leq b$を整数として$I=B(a,b-a+1)$と置きます.

$I=\int_0^1x^{a-1}(1-x)^{b-a}dx=\int_0^1\sum_{i=0}^{b-a}(-1)^{i}\binom{b-a}{i} x^{a+i-1}dx$

$=\sum_{i=0}^{b-a}(-1)^{i}\binom{b-a}{i}\frac{1}{a+i}$

最後の結果の分母に現れる整数は$a,...,b$で,被積分関数が常に正なので$d_b I$は正の整数になる.
また,$I$をガンマ関数で書き直すと

$I=\frac{\Gamma(a)\Gamma(b-a+1)}{\Gamma(b+1)}=\frac{(a-1)!(b-a)!}{b!}=\frac{1}{a}\frac{a!(b-a)!}{b!}$

$=\frac{1}{a\binom{b}{a}}$

ここで$d_b I$は正の整数より$a\binom{b}{a}|d_b$となることが分かります.
ここまでの議論に対し$(a,b)=(n,2n),(n+1,2n+1)$ $(n\in\mathbb{Z}^+)$を代入すると

$n \binom{2n}{n} \mid d_{2n},(2n+1)\binom{2n}{n}=(n+1)\binom{2n+1}{n+1}|d_{2n+1}$

が成り立ちます.
ここで明らかに$d_{2n}|d_{2n+1}$なので上で登場した$ n\binom{2n}{n},(2n+1)\binom{2n}{n}$は$d_{2n+1}$の約数であると分かります.
また$n,2n+1$は互いに素なので$n(2n+1)\binom{2n}{n}|d_{2n+1}$なので

$n(2n+1)\binom{2n}{n}\leq d_{2n+1}$

が言えます.

$(1+1)^n$の二項展開を考えた時最大の項は$\binom{2n}{n}$で項数は$2n+1$個なので

$(2n+1)\binom{2n}{n}\geq 2^{2n}=4^n$

と示せました.
このこと先ほどの議論から

$d_{2n+1}\geq n4^n$

よって$n\geq 2$では

$d_{2n+1}\geq 2^{2n+2}$

$n\geq 4$では

$d_{2n+2}\geq d_{2n+1}\geq 2^{2n+2}$

これらの不等式から$m=7,9,m\geq 10$では$d_m\geq2^m$であることが示せた.具体的に$m=8$の時は$d_8=840\geq2^8=256$より成り立っているため補題が示せた.

この補題を用いて次のLemma 2を示していきます.先ほどのSTEP2のコードを載せた時に紹介すると言った証明です.

Lemma 2

$Ord_{r}(n)>{(log_2 n)^{2}}$を満たす$r$が$max[3,{(log_{2} n)}^{5}]$までに存在する

(proof)
任意の正の整数$a$について

$\prod^a_{i=1}n^i=n^{a(a+1)/2}$

であり,$a\geq4$の条件下では$1+\frac{a(a+1)}{2}<a^2$であり,$\lfloor(\log_2 n)^2\rfloor\geq4$なので

$M=\prod^{\lfloor(\log_2 n)^2\rfloor}_{i=1}(n^i-1)$

と定義してあげると,$n$と$M$は互いに素なので

$nM<n・n^{\lfloor(\log_2 n)^2\rfloor(\lfloor(\log_2 n)^2\rfloor+1)/2}<n^{(\log_2 n)^4}=2^{(\log_2 n)^5}$

が成り立つことがわかる.ここで主張に戻ると,まず$n=2$の時は$r=3$より明らかに成り立つので以降は$n>2$として考えていきます.
${(log_{2} 3)}^{5}$を計算すると$10$より大きいことがわかるので,${(log_{2} n)}^{5}>10$であり,さらにLemma 1を用いると

$d_{{(log_{2} n)}^{5}+1}\geq2^{{(log_{2} n)}^{5}}$

が言えます.このことから

$1<s\leq\lfloor(\log_2 n)^5\rfloor+1$

の範囲で$nM$の約数でないものが存在します.

$s_1=\frac{s}{gcd(s,n)}$　とおくと

$s_1\leq\lfloor(\log_2 n)^5\rfloor+1$

であり,定義の仕方から$s_1$と$n$は互いに素で$s$が$nM$の約数でないため$s_1$は$M$の約数ではないと分かります.したがって

$n^i\not\equiv1 \pmod{s_i}$ $(i=1,...,\lfloor(\log_2 n)^2\rfloor)$

であるため,

$Ord_{s_i}\geq\lfloor(\log_2 n)^2\rfloor+1>(\log_2 n)^2$

よって題意は示せたと分かります.

ここからはこのアルゴリズムで漏れなく素数を判定できるのか証明していきます.
ここでアルゴリズムの過程を再掲します.

アルゴリズムの流れ

まずはこのアルゴリズムの流れを紹介します。ここで判定したい整数を$n$とする.

$n$が整数の冪乗か調べる.(冪乗なら合成数のため終了)
$Ord_{r}(n)>{(log_2 n)}^{2}$を満たす最小の$r$を求める.
$1\leqq a\leqq r$に対して$a$と$n$が互いに素か調べる(そうでないなら合成数のため終了)
$n\leqq r$なら素数.そうでなければ次に行く
$a=1$から$a=\lfloor \sqrt{\phi (r)}log_2 n\rfloor$までの$a$について${(x+a)}^{n}\equiv x^{n}+a\pmod{x^{r}-1,n}$が成り立てば素数

まずはnが素数の時正しく素数であることを返すかを確認します.
STEP1,STEP3,STEP4で合成数であると判定されることはないため,STEP5を確認していきます.
$n$が素数より任意の$a$において$gcd(a,n)=1$よりフェルマーの小定理から${(x+a)}^{n}\equiv x^{n}+a\pmod{n}$が成り立つとわかります.よってSTEP5の範囲の$a$においても同様に成り立つため${(x+a)}^{n}\equiv x^{n}+a\pmod{x^{r}-1,n}$が成り立つことがわかります.よって素数はきちんと素数と返してくれることがわかります.

次に$n$が合成数のとき,正しく合成数であると返してくれるかどうかを考えましょう.
この仮定によって$n\geq4$であるのでこの元で考えていきます.

(i)nが偶数の時

このときSTEP2で求められる$r$が必ず2以上になるためSTEP3で$a$と$n$が互いに素か調べる段階で合成数と判断されることが分かります.

(ii)nが奇数の時

$n$が整数の冪乗だった場合STEP1で弾かれるため以降はそうでないと仮定します.
$Ord_r(n)>1$なので$n$の素因数$p$で$Ord_r>1$となるものが存在します.$p\leq r$ならば
STEP3で合成数と判断されるので$p>r$とさらに仮定します.ここで$p,n$は$r$と互いに素です.以降$p,r$を固定して

$s=\lfloor \sqrt{\phi (r)}\log_2 n\rfloor$

として考えていきます.
これ以降のSTEPで合成数が素数と判断されてしまう可能性が残っているのはSTEP5なのでここについて考察していきましょう.
そこで$a=1$から$a=s$までの$a$について${(x+a)}^{n}\equiv x^{n}+a\pmod{x^{r}-1,n}$が成り立つと仮定して矛盾を導いていきます.

4.AKSアルゴリズムがクラスPに属することの証明

素数を判定できることがわかったので次にクラスPに属していることを証明していきます.

記法と前提

論文では以下のような記法が用いられています.
■ $\tilde{O}(t(n))$ 記法

$\tilde{O}(t(n)) = O\left(t(n)\cdot \mathrm{poly}(\log t(n))\right)$
という略記。
任意の $\varepsilon > 0$ について
$\tilde{O}(t(n)) = O(t(n)^{1+\varepsilon})$
と同程度とみなせます.

■ 基本演算の時間

m ビット整数の加算・乗算・除算はすべて
$\tilde{O}(m)$
で実行可能と仮定します。
この証明はとても大変なので今回は仮定させてください.

AKSアルゴリズム各ステップの計算量解析

以下で $r$ は
$\operatorname{ord}_r(n) > \log^2 n$
を満たす最小の整数であり,先ほど示したLemma 2から
$r \le \lceil \log^5 n \rceil$
が保証されています.

STEP1：完全冪のチェック

$n = a^b \ (b>1)$ かどうかを調べる。
時間
$\tilde{O}((\log n)^3)$

STEP2：条件を満たすrの探索

$\operatorname{ord}_r(n)>(\log n)^2$となる最小の r を見つける。

試行回数：$O((\log n)^5)$

各 $r$ に対して $O((\log n)^2)$ 回の合同式チェック

1 回のチェック時間：$\tilde{O}((\log n)^2 \log r)$

各$r$に対してこの計算時間がかかるので総計算時間は単純に積を考えて

総時間:$O((\log n)^5)\cdot \tilde{O}((\log n)^2 \log r)= \tilde{O}((\log n)^7)$

STEP3・4

$gcd$を $r$ 回計算するので

$O(r\log n)=O(\log^6 n)$

また$n \le r$の判定については

$O(\log n)$

STEP5：多項式合同式の検証（最も重い部分）

本質的にはここの計算時間が一番大事です.
次の合同を L 回チェックする：

$(X+a)^n \equiv X^n + a \pmod{(X^r-1,\ n)}$

反復回数は
$L = \left\lfloor \sqrt{\varphi(r)} \log n \right\rfloor$
であり,1 回のチェックには次数$r$の多項式の乗算を$O(\log n)$回行う。

よって 1 回あたりの時間は
$\tilde{O}(r\log^2 n)$

総計算時間の計算

したがって総時間は
$\tilde{O}(L\cdot r\log^2 n)
= \tilde{O}\left(r\sqrt{\varphi(r)}\log^3 n\right)$

$\varphi(r) < r$より
$\tilde{O}(r^{3/2}\log^3 n)$

そして $r=O(\log^5n)$を代入すると
$\tilde{O}\left(((\log n)^5)^{3/2}(\log n)^3 \right)
= \tilde{O}((\log n)^{15/2} \cdot (\log n)^3)
= \tilde{O}((\log n)^{21/2})$

これは$\log n$の多項式時間なので,AKSアルゴリズムは決定的かつ無条件にクラスPに属すると分かりました.

5.最後に

今回はAKSアルゴリズムの実装と証明をしてきました.素数判定はとても奥が深く,様々な種類があります.例えばリーマン予想を仮定したMillerテストやそれを修正したMiller–Rabin素数判定法など他にも興味深いものがたくさんあります.また素数判定まわりでまだ証明されていないこともたくさんあるので,興味を持った方は是非調べてみてください.

6.参考文献

[1]M. Agrawal, N. Kayal, N. Saxena. PRIMES is in P,2003
[2]雪江明彦,整数論3 解析的整数論への誘い,日本評論社,2021
[3]素数判定いろいろ - AKS素数判定法