ACL (AtCoder Library) の内部実装を知りたい人向けの記事です。
お友だちに「ねーね、ACL のこの関数ってどういう仕組みなの？ 知ってたりしない？」と聞かれたとき、「え... なんかほら、わかんないけど、魔法で動くからいいんだよ」としか言えないと情けない気がしません？
しました。なので書きます。
こういうシチュエーションはなくても、何かを実装したいときに「あ、これ ACL の実装のやつと同じ発想じゃん」となることはありえるので、知っていて損はないかなと思います。
めちゃくちゃ長くなったので、一度に全部読むのには適さないかもしれません。
数式部分の LaTeX コードなども含めて数えられていますが、26000 文字を超えています。
全体像
github.com
math.hpp にある関数たちです。
#include "atcoder/internal_math"
namespace atcoder {
long long pow_mod (long long x, long long n, int m);
long long inv_mod (long long x, long long m);
std ::pair <long long , long long > crt (const std ::vector <long long >& r,
const std ::vector <long long >& m);
long long floor_sum (long long n, long long m, long long a, long long b);
}
github.com
atcoder /internal_math（を経由して #include される atcoder /internal_math.hpp）にある関数たちです。
namespace atcoder {
namespace internal {
constexpr long long safe_mod (long long x, long long m);
struct barrett {
unsigned int _m;
unsigned long long im;
barrett (unsigned int m);
unsigned int umod () const ;
unsigned int mul (unsigned int a, unsigned int b) const ;
};
constexpr long long pow_mod_constexpr (long long x, long long n, int m);
constexpr bool is_prime_constexpr (int n);
template <int n>
constexpr bool is_prime = is_prime_constexpr (n);
constexpr std ::pair <long long , long long > inv_gcd (long long a, long long b);
constexpr int primitive_root_constexpr (int m);
template <int m>
constexpr int primitive_root = primitive_root_constexpr (m);
unsigned long long floor_sum_unsigned (unsigned long long n,
unsigned long long m,
unsigned long long a,
unsigned long long b)
}
}
よくある実装と比べて、高速化などの工夫がされています。
これを読んでいる人がびっくりするであろう点を挙げます。
is_prime_constexpr の最悪計算量が \(\Theta(\sqrt{n})\) ではなく \(\Theta(\log(n))\)。
この関数はコンパイル 時だけでなく実行時にも呼べます。
pow_mod において、除算・剰余算を高々 2 回しか行わない。
crt において、計算途中でオーバーフローしない。
法たちの最小公倍数がオーバーフローしない前提です。
下手にやるとオーバーフローすると思います。
また、人によっては実装自体が思いつかない場合がありそうなものもあります。
inv_mod
法が素数 であることを仮定する実装しか知らない人もいそうです。
↑ Fermat の小定理により \(m-2\) 乗するやつです。
floor_sum
ACL 登場前にこれを求める問題が出ていたら、500, 600 点くらいはつきそうです。
primitive_root_constexpr
これらを当然に感じる人にとっては、目新しいことは書いていないかもしれません。
\[
\gdef\M{m_{\text{inv}}}
\]
個別の説明
説明をしていきます。
「この操作は非自明ですが後述する〇〇によってできます」のような説明が増えると、その時点で何が非自明かわからなくなりそうなので、予め武器は揃えておきたいです。
そういう理由で、internal_math の方から説明していきます。
（追記 2021/01/22：これらはあくまで math の内部利用のための関数で、ドキュメントには含まれていません。
なので、今後のバージョンにおいては実装が変わることもありえると思います。
たとえば、未来のバージョンにおいて atcoder::internal::is_prime_constexpr を呼んだら、ここで書かれていた実装と違って困った（or 存在しなかった）ということはありえそうな気はします。
適宜、この記事の時点でのバージョン などを GitHub で探して、コピペしたりするのもいいと思います。ライセンスが CC0 なので大丈夫だと思います。）
internal::safe_mod\(x\) と \(m \ge 1\) を受け取り、\(x \bmod m\) を返します。
これはほぼ自明なんですが、C++ の除算の仕様で、\(x \lt 0\) の場合には注意が必要です。
internal::barrettはい出ました。
法 \(1\le m\lt 2^{31}\) と \(0\le a\lt m\), \(0\le b\lt m\) に対して \(a\cdot b\bmod m\) を高速に求めるためのクラスです。
まず前提として、除算は他の演算（加算や乗算、ビット演算など）と比べて重いです。
なるべくなら他の演算に置き換えてしまいたいです。
他の演算に変換する方法はいくつか知られており、その一つがこの Barrett reduction です*1 。
実は、コンパイル 時定数であれば、この変換（最適化）はコンパイラ ちゃんがしてくれます（えらい）。
（実際には、いくつかの種類から状況に応じて効率のよいものをやってくれます。
実装は この辺 に書いてある気がします。あんまりちゃんと読めていませんが、ここで説明するものとは違うことをしているかもしれません。）
long long MOD = 998244353 ;
int main () { }
よりも
long long const MOD = 998244353 ;
int main () { }
の方が速くなったという経験があるという人もいると思いますが、これはそういう事情によるものです。
グローバルに const でない MOD が置かれていると、コンパイラ ちゃん的にはそれが定数なのか判断しかねるので、最適化してくれません*2 。
const をつけることで除算の部分が最適化され、速くなったわけですね。
コンパイラ ちゃんが最適化してくれるのはコンパイル 時定数だけなので、実行時に mod が決まるような場合にはつらいです。
前置きが長めになりましたが、この変換を実行時に自前でやっちゃおうねというのがこのクラスです。
Barrett reduction の話
Wikipedia が結構わかりやすいですが、ACL のコード中にコメントもあるので、それに従って説明していきます。
Barrett reduction は、法 \(m\) と \(0\le z \lt m^2\) に対して \(z\bmod m\) を求める方法です。
\(m\) に対する前処理で一度除算を行うのみで、各クエリ \(z\) に対しては 除算を行わずに 求めることができます。
\(0\le a\lt m\) と \(0\le b\lt m\) に対して \(0\le a\cdot b \lt m^2\) ですから、これを用いることで、乗算 \(a\cdot b\bmod m\) を除算なしで求めることができるわけです。
以下、\(z = a\cdot b \lt m^2\) ということにして話を進めます。
求めたいのは \(z \bmod m\) ですが、これは \(z - \lfloor z/m\rfloor \cdot m\) と等しいので、\(\lfloor z/m\rfloor\) を考えます。
基本的な方針は次のような感じです。
前計算で \(2^{64}/m\) くらいの値 \(\M\) を求めておく。
クエリに対して \(z\cdot \M\) を計算すると、\(z\cdot 2^{64}/m\) くらいの値になる
なので、\(\lfloor z\cdot \M/2^{64}\rfloor\) は \(z/m\) くらいの値になる
誤差を修正する
さっきから「くらい」と言っている部分のこと
この誤差は、これから証明するように高々 \(1\) である
→ 正しい値ならそのまま、そうでなければ \(1\) ぶん補正すればよい
3 の除算について補足（クリックで展開）
「64 bit の整数 \(z\), \(\M\) の積の、上位 64 bit（1 ワードぶん）」が欲しいという状況です。
実装の際は、以下のようにすればよいです。
((unsigned __int128)z * m_inv) >> 64
シフト演算は軽そうですが 128 bit 整数の乗算はどうなの？ という声も聞こえます。
実際には、コンパイラ ちゃんは賢くて、「64 bit 整数同士の積の上位ワードが欲しいのね」というのをわかっているような最適化をしてくれるので、助かります。
unsigned __int128 が使えない環境もあるんですが、同様のことをする命令 _umul128 を使うことで対処していますね。
See here .
3 の除算についての話おわり。
さて、ACL では、\(\M = \lceil 2^{64}/m\rceil\) としているので、以下そうします。
切り上げ除算は (a+b-1)/b がよく知られていますが、次のようにはできません。
m_inv = ((1ull << 64 ) + m-1 ) / m;
1ull << 64 の時点で未定義動作になってしまうので。(((unsigned __int128)1 << 64) + m-1) / m とかするといいのかも？ 重そうかも？
ということで、少し頭をひねります。
m_inv = ((1ull << 64 ) - 1 ) / m + 1 ;
m_inv = (-1ull ) / m + 1 ;
符号なし整数はオーバーフローした場合の動作が定義されている（循環する (wrapping)）ので、こういうことができるんですね。
Rust とかで実装する場合は気をつけましょう。
特に、m == 1 の場合は -1ull + 1 なので、そこの wrapping にも注意です（0 になってくれればよい）。
正当性の話
誤差が高々 \(1\) であることと、その検出方法を示して終わりです。
コード 中にもコメントがあります。
まず \(m = 1\) の場合、\(a = b = 0\) なので、\(z = 0\) です。また、\(\M = 0\) です。
\(z\cdot\M/2^{64} = 0\) なので、誤差はなくて、よいです。
残りは \(m \gt 1\) の場合です。
\(m\cdot\M = 2^{64} + r\) (\(0\le r \lt m\)) と書けます。\(r\) の範囲は天井関数の性質からきています。
また \(z \lt m^2\) なので、ある \(0\le c, d\lt m\) が存在して \(z = cm+d\) と書けます。
このあまり \(d\) が最終目標ですが、商 \(c\) に注目します。
商の方が求めやすくて、正しい商を求めればあまりは求められるので。
\[
\begin{aligned}
z\cdot \M
&= (c\cdot m+d)\cdot \M \\
&= c\cdot (m\cdot\M) + d\cdot \M \\
&= c\cdot (2^{64}+r) + d\cdot \M\\
&= c\cdot 2^{64} + c\cdot r + d\cdot \M.
\end{aligned}
\]
この \(c\cdot r + d\cdot \M\) について見ていきます。
\[
\begin{aligned}
c\cdot r + d\cdot \M
&\lt m\cdot m + m\cdot \M \\
&\lt m\cdot m + 2^{64} + m \\
&= 2^{64} + m\cdot(m+1) \\
&\lt 2^{64} + 2^{64} \\
&= 2\cdot 2^{64}.
\end{aligned}
\]
これにより、結局次のことがわかります。
\[
c\cdot 2^{64} \le z\cdot \M \lt (c+2)\cdot 2^{64}.
\]
要するに、\(\lfloor z\cdot \M/2^{64}\rfloor \in \{c, c+1\}\) ということです。
長いので、\(v = \lfloor z \cdot \M/2^{64}\rfloor\) としておきます。
\(v = c\) ならば \(z-v\cdot m = z-c\cdot m = (z\bmod m)\) で、\(0\le z-v\cdot m\lt m\) が成り立ちます。
\(v = c+1\) ならば \(z-v\cdot m = z-(c+1)\cdot m = (z\bmod m)-m\) で、\(-m \le z-v\cdot m\lt 0\) が成り立ちます。
よって、\(z-v\cdot m\lt 0\) のときに補正を行えばよいです。
長かったです、お疲れさまでした。
barrett のメンバ関数 の説明をしていきます。
barrett::barrett（コンストラク タ）では、上で言う \(\M\) を求めています。
barrett::umod は unsigned int 型で法 \(m\) を返すだけです。
barrett::mul は、上の方法で \(z\bmod m\) を求めます。
ちょっと横道の話
ここではあまりを求めるために Barrett reduction を使っていますが、途中で商を求めたことから分かる通り、除算の高速化としても使えます。
せっかくなので例でも見て遊びましょうか。
100, 3141592, 123456 などを 3 で割ってみましょう。
m_inv = (0xffffffffffffffff / 3 + 1);
// == 0x5555555555555556
(3 * m_inv) >> 64 // == 1
(3141592 * m_inv) >> 64 // == 1047197
3141592 / 3 // == 1047197
(123456 * m_inv) >> 64 // == 41152
123456 / 3 // == 41152
コンパイル 時定数であればコンパイラ ちゃんがやってくれますが、実行時に決まる値で何度も割り算しなきゃな状況では、こういう方法もあるということですね。
正直な話、この定数倍高速化によって TLE/AC が分かれるようなコンテストには出たくないですが。
fastest を取りたい人や、想定より重いオーダーの解法をがんばって通したい人向けくらいの話だと思います。
そういえば、rolling hash で実行時に法を決めたい人は、これを使うことで高速化が見込める気がします。
それでも、自前実装は必要ないはずで、ACL の dynamic_modint を使えばいいとは思います。
internal::pow_mod_constexpr\(x\), \(n \ge 0\), \(m \ge 1\) を受け取って \(x^n \bmod m\) を返します。
よくある繰り返し二乗法です。
\( (x, n, m) = (0, 0, 1)\) の場合にうっかり \(1\) を返さないように注意です。
ここ、完全に説明を省くつもりだったんですが、冷静になると知らない人もいるかもしれないので書きます。
愚直にやると \(\Theta(n)\) 時間かかりますが、これを \(O(\log(n))\) 時間で行います。
答えを \(r = 1\) で初期化して、これに適切な値を掛けることを繰り返します。
\(n\) を 2 進数で表したときの、下から \(i\) 桁目が \(1\) なら \(r\) に \(x^{2^i}\) を掛けます。
\(0\) なら何もしません（説明上は \(1\) を掛けると見なす方がいいかも）。
また、\(x^{2^i}\) が得られている状態では、\(x^{2^{i+1}} = x^{2^i} \cdot x^{2^i}\) と計算できます。
\(x^{2^0} = x\) とわかることから、順に計算していけます。
\(n\) の 2 進数での下から \(i\) 桁目は \(\lfloor n/2^i\rfloor \bmod 2\) として分かります。
実際には、処理ごとに \(n \gets \lfloor n/2\rfloor\) と更新していけば、\(n \bmod 2\) で判定できます。
（さらに実際には、n >>= 1 と n & 1 をしています。）
数式で確認しておいてみます。\(n\) がある \(k\), \(b_i \in\{0, 1\}\) (\(0\le i\lt k\)) に対して次のように表せるとします。
\[
n = \sum_{i=0}^k b_i\cdot 2^i.
\]
このとき、\(r\) は次のようになりますね。
\[
\begin{aligned}
r &= \prod_{i=0}^k x^{b_i\cdot 2^i} \\
&= x^{\sum_{i=0}^k b_i\cdot 2^i} \\
&= x^n.
\end{aligned}
\]
internal::is_prime_constexprはい出ました。
\(n \ge 0\) を受け取って、素数 なら true、そうでなければ false を返します。
よくある試し割り法だと最悪時 \(\Theta(\sqrt{n})\) 時間になってしまいますが、\(O(\log(n))\) 時間のアルゴリズム が用いられています*3 。
Miller–Rabin のアルゴリズム が関係していて、それを知らないことには無理なので、まずそれの話をします。
Miller–Rabin の話
\(3\) 以上の奇数 \(n\) を受け取り、それが素数 か判定するアルゴリズム です。
偶数または \(1\) の場合は簡単に分かるので、簡単にやってください。
大まかなアイディアは次のような感じです。
整数 \(a\) を \(a\not\equiv 0 \pmod{n}\) からランダムに選ぶ。
「これを満たさなければ合成数 である」知られている条件 \(P_a(n)\) を調べる。
残念ながら、「これを満たせば素数 である」とは限りません。
十分な個数の \(a\) を試し、それでも「合成数 である」とならなければ、それなりの確率で素数 でしょうというものです。
「それなりの確率」は適当ではなくて、個数に対して見積もられていますが、この記事の範囲外なのでスキップします。
ACL で使われているアルゴリズム では、この確率を \(1\) にする工夫がされているためです（次のお話が楽しみですね）。
さて、まずは \(P_a(n)\) について考えていきます。
う〜んう〜ん（考えている）。
Fermat の小定理から、素数 \(p\) と \(1\le a\lt p\) に対して \(a^{p-1}\equiv 1\pmod{p}\) が成り立ちます。
素数 \(p\) を法とする \(1\) の平方根 は \(1\) と \(p-1\) しかありません*4 。
よって、\(1\) でも \(p-1\) でもないものを \(2\) 乗して \(1\) になると、合成数 だとわかります。
奇数 \(d\) と整数 \(s \ge 1\) を用いて \(n = d\cdot 2^s+1\) と表します。
法 \(n\) の下で \(a^d, a^{d\cdot 2}, \dots, a^{d\cdot 2^s} = a^{n-1}\) を順に計算していきます。
このとき、\(1\) または \(p-1\) が現れたらそこで終了し、そのときの値を \(y=a^{d\cdot 2^i}\) とします。
（追記 2021/02/16：現れなければ最後まで計算して、\(y=a^{n-1}\) とします。）
\(i = 0\) で \(y\equiv \pm 1\) だった場合は許します。
\(i\gt 0\) で \(y \not\equiv p-1\) だった場合は「\(\pm 1\) 以外の平方根 が存在した*5 」か「\(a^{d\cdot 2^s}\) まで計算しても \(1\) にならなかった」ことになり、合成数 であることになります。
これを条件 \(P_a(n)\) とします。すなわち、\(P_a(n) = ( (i=0) \vee (y\equiv -1))\) です。\(P_a(n)\) が偽のとき、\(n\) は合成数 です。
頭が暗黙に素数 mod を仮定してしまって想像できない人*6 のために例を挙げてみます。
法 \(n\) を \(21 = 5\cdot 2^2+1\) とします。\(a = 8\) とすると \(a^d = 8^5 \equiv 8\)、\(a^{d\cdot 2} \equiv 1\) となります。
また、\(a = 3\) とすると \( (a^d, a^{d\cdot 2}, a^{d\cdot 2^2}) \equiv (12, 18, 9)\) となり、\(1\) になりません。
このような真偽判定の証拠に使われる \(a\) は証人 (witness) と呼ばれたりします。
証人の \(a\) ちゃん「自分は \(n\) が合成数 であることを証明できます」
...ということですね。
他にも、SAT（充足可能性問題）は一般に解くのが難しいですが、「この割り当てなら充足可能だよ」という証拠となる割り当てのことを witness と呼んだりしますね。
certificate とも呼ばれるかも。
確率的じゃなくする話
上記のアルゴリズム は基数 \(a\) を選ぶところで乱択を行っています。
\(n \le 2^{32}\) であれば、これは \(\{2, 7, 61\}\) のみを使えば十分であることが知られています。
なので、それをします。
\(n\in\{2, 7, 61\}\) なら素数 なのは知っているので、そのときはすぐ true を返しています。
\(a\equiv 0\) になってしまい、上記の判定法では false になってしまって都合が悪いためです。
\(7\) や \(61\) の倍数だった（かつ \(7\) や \(61\) でない）場合も怪しい感じはありますが、結果的には false になってくれるのでよいです。
また、入力に応じて適切に法を選択することで、1 つの \(a\) のみで判定する方法も知られています（提出コード (Rust) ）。
64-bit 整数においては、7 つの基数 \(\{2, 325, 9375, 28178, 450775, 9780504, 1795265022\}\) を用いる方法と、入力に応じて 3 つの基数を選ぶ方法が知られています。
オーバーヘッドが大きいものの、入力に応じて 2 つの基数を選ぶ方法も知られています。
7 つの基数の方は、合成数 も含まれているため、互いに素などの条件を気にしないとこわれる気がします。
それらの（論文の著者によると思われる）コードは ここ で公開されています。
internal::is_primeis_prime<173> のような書き方ができます。テンプレート定数ですね。
コンパイル 時に素数 判定が必要になるとき用のものです。
modint の法が素数 かどうかの判定に使われています。
internal::inv_gcd\(a\) と \(b \ge 1\) を受け取り、以下の条件を満たす \( (g, x)\) を返します。
\(g = \gcd(a, b)\), and
\(a\cdot x\equiv g \pmod{b}\), and
\(0\le x\lt b/g\).
\(a\) を適当に処理して、\(0\le a\lt b\) とします。
なお、\(a = 0\) の場合は \( (b, 0)\) を返しちゃいます。
互除法をベースにしつつ、適切な不変条件を管理しながらループします。
\(a\) と \(b\) の最大公約数 \(g = \gcd(a, b)\) について、\(x\), \(y\) が存在して次の等式が成り立ちます。
\[
ax + by = g.
\]
いま欲しいのはこの \(x\) で、\(y\) はいらないので、以下では \(b\) を法として考えます。
\[
\begin{aligned}
ax&\equiv g\pmod{b}; \\
g-x\cdot a&\equiv 0 \pmod{b}.
\end{aligned}
\]
互除法は知っていると仮定していいですか？
\( (a, b) \gets (b\bmod a, a)\) で更新していくと最終的に \( (0, g)\) になります。
それを踏まえつつ、上の式を目標として変形を繰り返します。
初期値は後述するとして、次の合同式 を管理します。
\[
\begin{aligned}
\begin{cases}
s - m_0\cdot a\equiv 0,\\
t - m_1\cdot a\equiv 0.
\end{cases}
\end{aligned}
\]
また、\(0\le x\lt b/g\) に関連して、次の不等式も考えておきます。
\[
s\cdot|m_1| + t\cdot|m_0| \le b.
\]
最終的に \( (s, t, m_0) = (g, 0, x)\) となってほしいのを考えると少し気持ちがわかるかも？しれません？
天下り かも。
\( (s, t) = (b, a)\) で初期化しておきます。
このとき、\( (m_0, m_1) = (0, 1)\) は合同式 と不等式を満たします。
\(s\gt t\) に注意しておきます。
さて、互除法を考えると \(s \bmod t\) に関する合同式 が欲しいです。
\[
(s\bmod t)-m_{\text{?}}\cdot a\equiv 0.
\]
ここで、\(s\bmod t = s-\lfloor s/t\rfloor\cdot t\) を思い出すと、\(s\) の式から \(t\) の式の \(\lfloor s/t\rfloor\) 倍を引きたくなります。
\[
(s\bmod t)-(m_0-m_1\cdot \lfloor s/t\rfloor)\cdot a\equiv 0.
\]
よって、\( (s, t) \gets (t, s\bmod t)\) と \( (m_0, m_1) \gets (m_1, m_0-m_1\cdot\lfloor s/t\rfloor)\) で置き換えるのを繰り返せば、以下の式が得られます。
\[
\begin{aligned}
\begin{cases}
g - m_0\cdot a\equiv 0,\\
0 - m_1\cdot a\equiv 0.
\end{cases}
\end{aligned}
\]
置き換えの際に不等式が保たれていることを確認しておきます。
\(s\cdot|m_1|+t\cdot|m_0|\le b\) を仮定します。
また、\(u = \lfloor s/t\rfloor\) とします。
\[
\begin{aligned}
&\phantom{{}\le}
(s-t\cdot u)\cdot|m_1|+t\cdot|m_0-m_1\cdot u| \\
&\le s\cdot|m_1|-t\cdot u\cdot|m_1| + t\cdot|m_0|+t\cdot|m_1|\cdot u \\
&= s\cdot|m_1| + t\cdot|m_0| \le b.
\end{aligned}
\]
\(u\) が非負であることと三角不等式をうまく使いましょう。
また、このときにオーバーフローしないことも確認しておきます。
\[
\begin{aligned}
|m_1\cdot\lfloor s/t\rfloor| &\le |m_1|\cdot s \\
&\le b, \\
|m_0-m_1\cdot\lfloor s/t\rfloor| &\le |m_0|\cdot t+|m_1|\cdot s \\
&\le b.
\end{aligned}
\]
さて、\(s = g \gt 0\), \(m_0 = x\) が得られたとしましょう。
この \(x\) の範囲について考えます。
ループ終了時に \(s = g\), \(t = 0\) になることから、その前のステップにおいて \(s \bmod t = 0\)（つまり \(s\) は \(t=g\) の倍数）だったことがわかります。
よって、適当な整数 \(k\), \(m_0\) を用いて次のように書けます（ここでの \(m_0\) は \(m_0=x\) のものではなく、前のステップのそれです）。
\[
\begin{cases}
kg - m_0\cdot a \equiv 0, \\
g - x\cdot a \equiv 0. \\
\end{cases}
\]
このとき、例の不等式に対して、次のようになります。
\[
kg\cdot|x|+g\cdot|m_0| \le b.
\]
\(g\cdot|m_0|\ge 0\) なので、
\[
kg\cdot|x| \le b.
\]
それから、元々 \(b\gt a\) だったので、互除法の過程から \(kg\gt g\) です（\(kg=g\) なら \(b=a\) じゃないとやばくない？）。
\[
g\cdot|x| \lt kg\cdot|x| \le b.
\]
長かったですが、結局、次のことがわかります。
\[
|x|\lt b/g.
\]
以上より、\(-b/g\lt x\lt 0\) または \(0\le x\lt b/g\) とわかるので、所望の範囲 \(0\le x\lt b/g\) になるように（条件 \(ax\equiv g\) を満たしつつ）調整しておわりです。
\(x\lt 0\) なら \(x + b/g\)、\(x\ge 0\) なら \(x\) が所望のものです。
長かったですね。
internal::primitive_root_constexpr素数 \(m\) に対して、法 \(m\) での最小の原始根 \(g\) を返します。
整数 \(g\) が次の条件を満たすとき、\(g\) は法 \(n\) における原始根と言います：
\(\gcd(a, n)=1\) なる任意の \(a\) に対して、整数 \(k\) が存在して \(g^k\equiv a\pmod{n}\) が成り立つ。
原始根が存在する条件は、\(n\) が以下の形式で書けることと同値です。
\(n = 1, 2, 4\), or
\(n = p^k\) for some prime \(p\), or
\(n = 2p^k\) for some prime \(p\).
ACL では一般の場合ではなく、素数 の場合のみに限定しています。
方針自体は単純で、\(g = 2, 3, \dots\) に対して「\(g\) は原始根である？」を高速に試し、yes であるものが見つかれば返すことをしています。
以下、\(m\) を素数 として話します。
\(g\) が原始根であることは、\(\{g^0, g^1, \dots, g^{m -2}\} = \{1, 2, \dots, m -1\}\) であることと同値です。
Fermat の小定理から \(g^{m -1}\equiv 1\) ですが、指数部が \(1\) 以上 \(m -2\) 以下では \(1\equiv g^0\) にならないということですね。
判定の話
「\(g\) は原始根である？」を高速に行う方法を考えていきます。
直前に書いた話から、\(\Theta(m)\) 時間掛ければ簡単に判定できそうですが、もっと速くしてほしいです。
ある整数 \(k \lt m -1\) で初めて \(g^k \equiv 1\) になってしまった状況を考えます（すなわち、\(1\le k'\lt k\) について \(g^{k'}\not\equiv 1\) です）。
このとき、\(k\) は \(m -1\) の約数になっていることを示します。
\(k=1\) のときは自明なので除外して、\(k\gt 1\) とします。
背理法 を使ってみましょうか。\(k\) が \(m -1\) の約数でないことは整数 \(i\ge 0\), \(1\le k'\lt k\) を用いて \(m -1 = i\cdot k+k'\) と書けます。
\[
\begin{aligned}
g^{m -1}
&= g^{i\cdot k+k'} \\
&\equiv (g^k)^i \cdot g^{k'} \\
&\equiv 1^i \cdot g^{k'} \\
&\equiv g^{k'} \not\equiv 1.
\end{aligned}
\]
これは \(g^{m -1}\equiv 1\) に反しています。おわり。
ということで、\(m -1\) の各約数 \(d\lt m -1\) に対して \(g^d\not\equiv 1\) を試せばいいのですが、もう少しよくできます。
\(m -1 = p_1^{a_1} \cdots p_s^{a_s}\) と素因数分解 できたとき、\(1\le j\le s\) に対して \( (m -1)/p_j\) の形のものだけ試せばよいです。
\(m -1\) のある約数 \(d\lt m -1\) を考えます。
ある \(j\) に対して \(d\) の倍数であるような \( (m -1)/p_j\) が存在します*7 。
すなわち、\(d\cdot k = (m -1)/p_j\) と書けます。
いま、\(g^d\equiv 1\) が成り立つとすると、次のようになります。
\[
\begin{aligned}
g^{(m -1)/p_j}
&= g^{d\cdot k} \\
&= (g^d)^k \\
&\equiv 1^k = 1.
\end{aligned}
\]
よって、\( (m -1)/p_j\) の形のものさえ調べれば十分とわかりました。
ここまでをまとめると、次のような方法になります。
素因数分解 \(m -1 = p_1^{a_1}\cdots p_s^{a_s}\) を求める。\(g = 2, 3, \dots\) に対して、以下を試す。
ある \(j\) について \(g^{(m -1)/p_j} \equiv 1\) であれば no
そうでなければ yes で、その \(g\) が答えとなる
補足の話
ACL 中に \(m\) が素数 であることを要求する旨が書かれているため、それを仮定したことを書きました。
実際には、\(m -1\) の部分を Euler's totient function \(\phi(m)\) で置き換えることで似たような議論ができる箇所が多くあります。
Note：素数 \(m\) に対して \(\phi(m)=m -1\) です。
詳しくは ここ を見るとよいと思います。
計算量の話
整数 \(n\) の素因数の種類数を \(\omega(n)\) と書く*8 ことにし、\(r = \omega(m -1)\) とおきます。
各判定には \(O(r\cdot \log(m))\) かかります。
一般化された Riemann 予想の下では、\(g = O(r^4\cdot(\log(r)+1)^4\cdot \log(m)^2)\) であることが知られているようです。
また \(\omega(n) = O(\log(n)/\log(\log(n)))\) が成り立つようなので、\(g = O(\log(m)^6)\) です。
\(m -1\) の素因数分解 にかかる時間を \(F(m -1)\) とすると、全体の計算量は \(O(F(m -1)+\log(m)^8/\log(\log(m)))\) 時間です...？？
\(g\) はかなり loose な見積もりなはずなので、実際にはこんなにはかからないはずです...？
実際には \(F(m -1) = O(\sqrt{m})\) の部分がボトルネック となりそうです。
\(O(\sqrt{\sqrt{m}})\) の素因数分解 も知られています（こことか ）が、どうでしょうね。
\(g\) の見積もりについては この論文 にあります。
読めていないです。
internal::primitive_rootこちらもテンプレート定数です。
internal::floor_sum_unsigned（追記 2021/01/19：これ により追加されました）
\(m\) の制約を見落としていて恥ずかしいコメントをしてしまいました......
floor_sum の入力を非負整数に限ったバージョンです。
floor_sum については後述しますが、以下の値を求める関数です。
\[
f(n, m, a, b) = \sum_{i=0}^{n-1} \left\lfloor\frac{ai+b}{m}\right\rfloor.
\]
元々は \(a\) と \(b\) の範囲に制約があったのですが、それを取り払うために追加されました。
行っていることは、後で説明する（元々の）floor_sum と同じなので、ここでは省略します。
追記おわり。
これで internal_math は終わりです。
pow_modはい出ました。
barrett を用いることで、除算の回数を抑えます。
Barrett reduction を用いる以外は普通の繰り返し二乗法です。
internal::pow_mod_constexpr では barrett を用いていませんが、この関数はコンパイル 時に呼ばれるのが主なはずだからかな？と思っています。
barrett は実行時の高速化用なので*9 。
internal::is_prime_constexpr も同じ事情だと思います。
inv_mod\(x\) と \(m \ge 1\) を受け取り、\(x^{-1}\) が存在すると仮定してそれを返します。
\(x^{-1}\) が存在する条件は \(\gcd(x, m) = 1\) です。
これが成り立たない場合は、assert(false) です*10 。
internal::inv_gcd を使います。
これの返り値が（仮定の下で）\( (1, x^{-1})\) です。
crtある未知の整数 \(y\) があります。
条件 \(y\equiv r_i\pmod{m_i}\) を満たす \( (r_i, m_i)\) (\(1\le i\le n\)) を与えます。
このとき \(m = \lcm_i m_i\) として、\(y\equiv r \pmod{m}\) なる \(r\) を探し、\( (r, m)\) を返します。
存在しなければその旨を返します。\(O(n)\) 時間です。
こうした \(0\le r\lt m\) が一意であることが中国剰余定理 (Chinese Remainder Theorem; CRT) から言えます。
元々の CRT は「\(m_0\), \(m_1\) が互いに素ならば \(r\) が一意に存在する」と言っていると思います。
ここでは、互いに素の制約が無く、存在判定も行う状況を考えます*11 。
\(n = 2\) の場合のみ考えます。
一般の場合には、\( (r_1, m_1), (r_2, m_2)\) に対する答えを \( (r, m)\) として \( (r, m), (r_3, m_3)\) を求め、... というのを順に繰り返していけばよいです。
ACL の実装では \( (r_0, m_0) = (0, 1)\) とおいてループしていますね*12 。
雑に実装すると、計算途中の値が大きくなってオーバーフローしがちなんですが、\(m=\lcm_i m_i\) を超えないように注意がなされています。
なので、それに気をつけながら説明していきます。
コード 中にコメントがあるので、それを読んでもいいと思います。
解の構成の話
さて、以下の説明では、\( (r_0, m_0), (r_1, m_1)\) を受け取り \( (r, m)\) を返すものとして説明します。
適当な前処理により \(0\le r_0\lt m_0\)、\(0\le r_1\lt m_1\)、\(m_0\ge m_1\) としておきます。
まず、\(m_0\) が \(m_1\) の倍数である場合を考えます。
\(m = m_0\) であり、\(m_0\) に関する条件は自明に成り立ちます。
\(r_0\not\equiv r_1\pmod{m_1}\) なら解なしです。
そうでないとき、\( (r, m) = (r_0, m_0)\) です。
残りの場合を考えます。
求めたい値 \(r\) に対して次の式が成り立っていてほしいです（そういう話をしているので）。
\[
\begin{cases}
r\bmod m_0 = r_0, \\
r\bmod m_1 = r_1.
\end{cases}
\]
一つ目の式から、\(x\) が存在して \(r = x\cdot m_0+r_0\) と書けます。
これと二つ目の式から、次のことが言えます。
\[
(x\cdot m_0+r_0)\bmod m_1 = r_1.
\]
これを \(x\) について解くのが目標です。
\[
(x\cdot m_0) \bmod m_1 = (r_1 - r_0) \bmod m_1.
\]
右辺も \(m_1\) で割ったあまりとしている理由は、\(r_1-r_0\ge 0\) とは限らないためです。
気持ちとしては、\(m_0\) で両辺を割りたいです。
\(g = \gcd(m_0, m_1)\) として、\(m_0 = u_0 g\)、\(m_1 = u_1 g\) とします。
\[
(x\cdot u_0 g)\bmod u_1 g = (r_1-r_0) \bmod u_1 g.
\]
左辺は \(g\) の倍数であることがわかるので、右辺も \(g\) の倍数である必要があります。
すなわち、\(r_1\not\equiv r_0\pmod{g}\) なら解なしです。
\(r_1\equiv r_0\) なら \(r_1-r_0\) は \(g\) で割れるので、次のようになります。
\[
(x\cdot u_0)\bmod u_1 = \left(\frac{r_1-r_0}{g}\right) \bmod u_1.
\]
右辺は、「\( (r_1-r_0)/g\) を \(u_1\) で割ったあまり」であって「\(r_1-r_0\) を法 \(u_1\) の下で \(g\) で割ったもの」ではないことに注意してください。
そもそも \(g^{-1}\pmod{u_1}\) は存在するとは限らないので。
ここまできたら、もう合同式 にしちゃっていいです。
\[
x\cdot u_0 \equiv \left(\frac{r_1-r_0}{g}\right) \pmod{u_1}. \\
\]
\(u_0\) と \(u_1\) は互いに素なので \(u_0^{-1}\pmod{u_1}\) は存在します。
実装の際は、\(g = \gcd(m_0, m_1)\) を求める際に internal::inv_gcd を使うことで \( (g, u_0^{-1})\) をまとめて得られます。
\[
x \equiv \left(\frac{r_1-r_0}{g}\right)\cdot u_0^{-1} \pmod{u_1}.
\]
この \(x\) のうち \(0\le x\lt u_1\) のものを使って \(r = x\cdot m_0+r_0\) を計算すればよいです。
また、\(u_1 = m_1/g\) より、\(m = m_0\cdot u_1\) です。
値の範囲の話
ここまでは計算方法の話をしていました。
あとは、絶対値が \(m = \lcm(m_0, m_1)\) 以下の値の演算で済ませられることを示します。
これにより、\(m\) がオーバーフローしない限り、オーバーフローが起きないことが言えます。
実際に計算する必要がある値は次の通りです。
\(r_1-r_0\),
\(x = (r_1-r_0)/g\cdot u_0^{-1} \bmod u_1\),
\(r = x\cdot m_0+r_0\),
\(m = m_0\cdot u_1\).
いま \(m_0\gt m_1\ge 1\) であり、以下の式が成り立つことに注意しておきます。
\[
\lcm(m_0, m_1)\ge 2\cdot\max\{m_0, m_1\}.
\]
三角不等式とか、この大小関係とかを使います。
\[
\begin{aligned}
|r_1-r_0| &\le |r_1|+|r_0| \\
&\lt m_1 + m_0 \\
&\lt 2\cdot\max\{m_0, m_1\} \\
&\le \lcm(m_0, m_1) = m.
\end{aligned}
\]
\(|r_1-r_0|\lt m\) より \(|r_1-r_0|/g\lt m\) はすぐわかります。
また、次が成り立ちます（掛け算の際にいちいち \(\bmod u_1\) してもよいといういつもの）。
\[
\left(\frac{r_1-r_0}{g}\right)\cdot u_0^{-1} \bmod u_1
= \left(\left(\frac{r_1-r_0}{g}\right)\bmod u_1\right) \cdot u_0^{-1} \bmod u_1
\]
\( ( (r_1-r_0)/g)\bmod u_1\) と \(u_0^{-1}\) はそれぞれ \(u_1\) 未満なので、\(u_1^2\lt m\) を示せば十分です。
\[
\begin{aligned}
u_1^2 &= (m_1/g)^2 \\
&\lt (m_0\cdot m_1)/g^2 \\
&\le (m_0\cdot m_1)/g \\
&= \lcm(m_0, m_1) = m.
\end{aligned}
\]
次いきます。
\[
\begin{aligned}
|r| &\le |r_0| + |m_0\cdot x| \\
&\lt m_0 + m_0\cdot |x| \\
&\le m_0 + m_0\cdot (u_1-1) \\
&= m_0 + m_0\cdot u_1 - m_0 \\
&= m.
\end{aligned}
\]
\(m = m_0\cdot u_1\) は特にいいですね。
正整数同士の乗算なので、右辺の各々より左辺は小さくならないです。
以上により、示したいことたちが示せました。
floor_sum以下のものを求めます。
\[
f(n, m, a, b) = \sum_{i=0}^{n-1} \left\lfloor \frac{ai+b}{m}\right\rfloor.
\]
詳しい解説は ここ に書きました。
大まかな流れは次のような感じです。
まず、\(a\gets a\bmod m\), \(b\gets b\bmod m\) とし、\(0\le a, b\lt m\) の場合に帰着させます。
このとき、\(0\le \lfloor (ai+b)/m\rfloor\le \lfloor(an+b)/m\rfloor\) であり、いろいろすると、次のような形に帰着できます。
\[
\sum_{i=0}^{\lfloor(an+b)/m\rfloor-1} \left\lfloor \frac{a'i+b'}{m'}\right\rfloor.
\]
この \( (a', m')\) が \( (m, a\bmod m)\) であり、\(f(n, m, a, b)\) が \(f(\lfloor(an+b)/m\rfloor, a\bmod m, m, b')\) に帰着されます。
\(m\) と \(a\) に対して互除法のようになっているので、それに関する計算量で抑えられます。
\(n\) も \(m\) や \(a\) と共に減っていき、\(O(\log(\min\{a, m, n\}))\) くらいの計算量です。
（追記 2021/01/19：\(a\), \(b\) の制約がなくなった件について補足）
上のリンク先で \(0\le a\lt m\) などに帰着させるパートがありますが、\(a\lt 0\) の場合も同様に帰着できます。
ただし、C++ の除算の仕様に関して注意が必要です。
internal::safe_mod を用いることで、\(\lfloor a/m\rfloor\) を \( (a-(a\bmod m))/m\) として計算しています。
余談
実は、ACL の floor_sum にはえ びちゃんが 関与 しています（うれしいね）。
Improve floor_sum [Merged]
おわり
おわりです。お疲れさまでした。
