O(1) mod inv, mod pow

概要
$O(1)$ inv
$O(1)$ log
1. 追記
2. 追記その 2（hos.lyric さんによる改善）
$O(1)$ pow
関連

概要

QOJ を問題番号順で眺めたり，典型トレーニングセットを眺めたことがある人は，次の問題を知っていると思います．

https://qoj.ac/contest/1167/problem/5

これは，必要ならば適当な事前計算を行うことで $x^{-1} \bmod p$ を高速に計算することを要求している問題で，問題名の示す通り，$O(1)$ 時間で online に計算することとができます．（もちろん前計算は $o(p)$ です）．

これを応用すると，mod pow を $O(1)$ 時間で online に計算することができることに気付いたので実装してみました．

以下 $p$ と書いたら奇素数です（$10^9$ 程度の大きさであることを想定しています）．

$O(1)$ inv

まず，証明します．

$n$ を正整数とする．任意の $x\in [1,p)$ に対して，$x \equiv a/b \pmod{p}$ かつ $0\leq |a| \leq 2p/n, 1\leq b\leq n$ となる $(a,b)$ が存在する．さらに $O(n^2)$ 時間の前計算によりこのような $(a,b)$ を $O(1)$ 時間で計算できる．

分母が $n$ 以下の有理数を並べた Farey 数列を考えることで，$c_1/b_1\leq x/p\leq c_2/b_2$ かつ $0\leq b_i,c_i\leq n$ となる $b_i,c_i$ であって，$b_1+b_2>n$ かつ $|b_1c_2-b_2c_1|=1$ となるものがとれます．

$b_1<b_2$ であるとき $b_2> n/2$ なので，$(b,c)=(b_1,c_1)$ とすると，

$$\left|\frac{x}{p}-\frac{c}{b}\right|\leq \left|\frac{c_1}{b_1}-\frac{c_2}{b_2}\right|=\frac{1}{b_1b_2}=\frac{2}{nb}$$

より $|bx-cp|\leq \dfrac{2p}{n}$ です．$b_2<b_1$ の場合も $(b,c)=(b_2,c_2)$ とすれば同じ不等式が得られます．$a=bx-cp$ とすれば求める $(a,b)$ が得られます．

ここまでの証明から，$(a,b)$ を求めるには，$x/p$ を分子・分母 $n$ 以下の有理数による左側・右側の最良近似分数が得られればよいことが分かります．

分子・分母 $n$ 以下の有理数が $n^2$ 個以下しかないので，これらすべてに対して何らかの処理を行えます．またそのような $2$ つの有理数の差は $1/n^2$ 以上なので，$[0,1]$ を $n^2$ 個のバケットに分割し，各バケットに対して前後の有理数を記録しておけばクエリに $O(1)$ 時間で答えることができます．

これを踏まえて，逆元を $x^{-1} = (a/b)^{-1} = b\cdot a^{-1}$ と計算することにします．$|a|\leq 2p/n$ という範囲の $a$ に対する逆元を前計算しておけば，逆元クエリが online $O(1)$ 時間になります．前計算の時間・空間は $O(n^2 + 2p/n)$ で，$n=p^{1/3}$ ととれば $O(p^{2/3})$ です．

QOJ への提出では全く同じコードの実行時間が 2 倍以上ずれたことがある（submissions 914536, 914538）ので，計測したいときには手元でやった方がいいと思います（インタラクティブだと不安定さが増幅するのだろうか）．

$O(1)$ log

以下 $n=p^{1/3}$ とします．

原始根 $r$ を固定して，対数写像 $\log \colon \mathbb{F}_p^\times \longrightarrow \mathbb{Z}/(p-1)\mathbb{Z}; r^x\longmapsto x$ を考えます．

逆元のときと同様，$1\leq |a| \leq 2p/n$ の範囲で $\log a$ が計算できていれば $\log$ が $O(1)$ 時間でクエリできます．$\log(ab)=\log(a)+\log(b)$ なので，$-1$ と素数に対する $\log$ だけ計算できればよいです．なお $\log(-1)=(p-1)/2$ です．素数定理で見積もれば，事前計算が必要な $\log$ の個数は，$2p/(n\log p)$ ということになります．

底を固定して $Q=2p/(n\log p)$ 個の離散対数を解く場合，BSGS 法のバケットサイズを適切にとることで，計算量を $O(\sqrt{pQ})$ 時間にできます．$n = p^{1/3}$ の設定でこれは $O(p^{5/6}/(\log p)^{1/2})$ 時間です．私の実装では $p=998244353$ 等で $0.75$ 秒程度で前計算ができることを確認しています．

Index Calculus も，底が固定されていて $Q$ クエリという条件下での高速化が見込める手法ですが，私は実装していません．

また，$\bmod 998244353$ の設定では $p-1$ が小さい素数の積であるため離散対数問題は高速に解けます．私の実装（Pohlig-Hellman algorithm の簡略化のような何か）では $0.1$ 秒程度ですべての前計算をできることを確認しています．

追記

次のような Index Calculus の簡略化のような方法で，$p$ の特殊性を使わずに構築 $0.17$ 秒を達成できました．

小さい $i$ については BSGS で解く．
大きい $i$ については，$k$ を乱択して $ir^k \bmod p$ を求める．これが $i$ 未満の素数の積に分解できれば，小さい場合の値を利用して $\log(i)$ が計算できる．素因数分解は十分小さい素数による試し割りと，ある程度以下の値に対しては篩を持っていることを利用して簡易的に行う．

詳細は後述の実装例を参照．

追記その 2（hos.lyric さんによる改善）

$i \geq \sqrt{p}$ のときの $\log(i)$ は，$p=ij+k$ として $\log(i) = \log(-k)-\log(j)$ を利用すると，$O(1)$ 時間で手に入ります．したがって，解くべき離散対数問題の個数は $O(p^{1/2}/\log p)$ 個となります．この方法で，構築 $0.1$ 秒以下が達成できました．

$O(1)$ pow

$a,b$ を与えて $a^b\pmod{p}$ を求める問題を考えます．$a=0$ のときを場合分けすれば，残りの場合は $a^b=r^{b\log a}$ により，$r$ のべき乗を計算する問題に帰着できます．

指数を $p-1$ で割った余りに置き換えて平方分割を考えれば，この部分については $O(\sqrt{p})$ 時間の前計算によりクエリ $O(1)$ 時間にできます．

$1\leq a<p$, $0\leq b<p-1$ のとき，私の実装ではクエリあたり，除算が $2$ 回と配列参照が $5$ 回行われています．より良い方法が見つかったら教えてもらえるとうれしいです．

実装例：https://maspypy.github.io/library/mod/modfast.hpp

手元でランダム $10^8$ クエリを処理した場合，

通常の繰り返し二乗法：$11.4$ sec
この記事の手法（前計算は除く）：$3.7$ sec

程度の高速化となりました．

概要

$O(1)$ inv

$O(1)$ log

追記

追記その 2（hos.lyric さんによる改善）

$O(1)$ pow

関連