AKS素数判定法をC++で実装する

はじめに

AKS素数判定法という素数判定アルゴリズムがあります。
詳しい内容については、該当しそうな論文や参考文献のリンク先などから調べていただきたいのですが、整数 $n$ が素数であるかどうかを、（入力ビット数が $\log_2{n}$ なので $\log{n}$ についての）多項式時間で判定できるアルゴリズムとなっています。よく見かける、時間計算量が $O\left(\sqrt{n}\right)$ の素数判定アルゴリズムは、入力サイズに対して見たときに指数時間のアルゴリズムとなっています。
多項式時間と聞くと、指数時間とされる $O\left(\sqrt{n}\right)$ の素数判定アルゴリズムより高速だと思えてしまうのですが、$n \sim 10^6$ 程度だと、AKS素数判定法の方が格段に遅いです。非常に大きな整数に対してはどのような関係になっているのか気になるところですが、手元では調べられませんでした。

実装

細かいアルゴリズムの詳細やステップに関しての説明は、参考文献に委ねます。（ステップに関してはWikipediaの記述に準拠していると思います。不等式は英語版のWikipediaに準拠していて、比較する値が原著論文とは少し違います。）
long long型やlong double型を多用していている割に、$N<2^{31}$ くらいの入力が限界だと思ってます。これは、途中で$\mod{n}$ の乗算を挟んでいるためです。
繰り返し二乗法を用いていたり、$r$ が $\log{n}$ の多項式で抑えられたりするので、実装にミスがなければ全体でも $\log{n}$ の多項式時間のアルゴリズムになっていると思います。

# include <cmath>
# include <iostream>
# include <numeric>
# include <vector>


namespace myAKS {
    // 配列形式で多項式を表す
    class polynomial_modulo {
        private:
            long long n, r, a;
            std::vector<long long> ls1, ls2;
            // a * x^n = a * x^(n%r) (mod x^r - 1)
            void set_ls1() {
                ls1[0] = a % n;
                ls1[n%r] = (++ls1[n%r]) % n;
            }
            void set_ls2() {
                std::vector<long long> tmp(r, 0);
                ls2[0] = 1;
                tmp[1] = 1; tmp[0] = a;  // x + a
                int n_ = n;
                while(n_ > 0) {
                    if(n_ & 1) {
                        ls2 = mul(ls2, tmp);
                    }
                    tmp = mul(tmp, tmp);
                    n_ >>= 1;
                }
            }
            std::vector<long long> mul(std::vector<long long> lhs, std::vector<long long> rhs) {
                std::vector<long long> retval(r, 0);
                for(long long i = 0; i < r; ++i) {
                    if(lhs[i] == 0) {
                        continue;
                    }
                    for(long long j = 0; j < r; ++j) {
                        retval[(i+j)%r] += lhs[i] * rhs[j] % n;
                        retval[(i+j)%r] %= n;
                    }
                }
                return retval;
            }
        public:
            polynomial_modulo(long long n_, long long r_, long long a_) {
                n = n_; r = r_; a = a_;
                ls1.resize(r, 0);  // x^n + a
                ls2.resize(r, 0);  // (x + a)^n
                set_ls1();
                set_ls2();
            }
            bool is_same() {
                for(long long i = 0; i < r; ++i) {
                    if(ls1[i] != ls2[i]) {
                        return false;
                    }
                }
                return true;
            }
    };

    // floor(log2(x)) の計算
    int floor_log2(long long n) {
        int log2_n = 0;
        n >>= 1;
        while(n > 0) {
            n >>= 1;
            ++log2_n;
        }
        return log2_n;
    }

    // 整数の累乗計算
    long long powint(long long n, int k) {
        long long tmp = n, retval = 1;
        while(k > 0) {
            if(k & 1) {
                retval *= tmp;
            }
            tmp *= tmp;
            k >>= 1;
        }
        return retval;
    }

    // 累乗数の判定
    bool is_perfect_power(long long n) {
        for(int i = 2; i <= floor_log2(n); ++i) {
            int root = roundl(pow(n, 1.0L / i));
            if(powint(root, i) == n) {
                return true;
            }
        }
        return false;
    }

    // 位数の計算
    long long order(long long n, long long mod) {
        long long mul = 1;
        for(long long e = 1; e < mod; ++e) {
            mul *= n;
            mul %= mod;
            if(mul == 1) {
                return e;
            }
        }
        return -1;
    }

    // 位数に関する不等式を満たすような最小の数を算出
    long long enough_order_modulo(long long n) {
        long long rhs = floorl(log2((long double)n) * log2((long double)n));
        for(long long r = 1; r < n; ++r) {
            if(order(n, r) > rhs) {
                return r;
            }
        }
        return n;
    }

    // 素因数列挙
    std::vector<long long> primefactors(long long n) {
        long long tmp = n;
        std::vector<long long> retval;
        for(long long i = 2; i*i <= n; ++i) {
            while(tmp % i == 0) {
                tmp /= i;
                retval.emplace_back(i);
            }
        }
        if(tmp > 1) {
            retval.emplace_back(tmp);
        }
        return retval;
    }

    // オイラーのトーシェント関数
    long long totient(long long n) {
        std::vector<long long> pfs = primefactors(n);
        long long retval = n;
        for(long long pf: pfs) {
            retval *= (pf-1);
            retval /= pf;
        }
        return retval;
    }

    // AKSアルゴリズムによる素数判定
    bool is_prime(long long n) {
        // 1(と0以下の整数)は素数でない
        if(n <= 1) {
            return false;
        }

        // Step1
        if(is_perfect_power(n)) {
            return false;
        }

        // Step2
        long long r = enough_order_modulo(n);

        // Step3
        for(long long a = 2; a <= std::min(r, n-1); ++a) {
            long long gcd_an = std::gcd(a, n);
            if(1 < gcd_an && gcd_an < n) {
                return false;
            }
        }

        // Step4
        if(n <= r) {
            return true;
        }

        // Step5
        long long rangemax = floorl(sqrt((long double)totient(r)) * log2l((long double)n));
        for(long long a = 1; a <= rangemax; ++a) {
            polynomial_modulo pm(n, r, a);
            if(!pm.is_same()) {
                return false;
            }
        }

        // Step6
        return true;
    }
}  // namespace myAKS


int main() {
    long long n;
    std::cin >> n;
    if(myAKS::is_prime(n)) std::cout << n << " is a prime number.";
    else std::cout << n << " is not a prime number.";
}

動作結果

C++17で動作を確認しました。
ARC017-AやARC044-Aなどに提出してみたのですが、AKSアルゴリズムは遅いという前評判通り、TLE（実行時間制限超過）となってしまいました。しかし、他はAC（正解）となっていたため、アルゴリズム自体はきちんと動作していると思われます。試しに、手元の環境で、TLEしてしまったARC017-Aの入力例3（$n=999983$）を実行してみると、1分程度かかってしまったものの、素数であるという正しい結果が得られました。
ただ、実装に全くミスがないとも言い切れないので、もし何か見つけた際には教えていただけると助かります。

参考文献

1. Wikipedia（日本語）
2. Wikipedia（英語）
3. 素数判定いろいろ - AKS素数判定法
4. 素数判定法