メモ

yukicoderでゆるふわgolf

Tonelli–Shanks algorithm

この記事ではTonelli–Shanks algorithmの原理の「気持ち」を説明します。
間違った説明はしていないつもりですが、(故意に)不正確な記述があるかもしれません。
ループ不変量に着目する考え方はmod_sqrtについて.md · GitHubを参考にしました。
(追記:上記サイトの更に参考元である英語版wikipediaにも同様の記述がありました)
アルゴリズムの説明のみであり、実装については一切説明していません。

Tonelli-Shanks algorithmとは次のような問題を解決するアルゴリズムである。
 問題:整数xと奇素数p、整数nが与えられた時、 g^2\equiv x \mod p^nをみたす整数gを高速に求めよ

より一般に次が解ける。
pが奇素数ならば、(\mathbb{Z}/p^n\mathbb{Z})^*は位数(p-1)p^{n-1}巡回群になることに注意せよ)
 問題:巡回群(G,+)とその元x\in Gが与えられた時、g+g=xを満たすg\in Gを高速に求めよ


以下、群の演算は全て加法的に書いているため、実際に乗法群(\mathbb{Z}/p\mathbb{Z})^*に対して用いる際は注意せよ


まずは簡単のためGの位数が2ベキの場合を考える
 問題:位数2^s (s>0)の巡回群(G,+)x \in Gが与えられる。g+g=xとなるg \in Gを求めよ。

\mathbb{Z}/2^s\mathbb{Z}の代表元を0~2^s-1とする。
 補題t \in Gが与えられた時、tの末尾に0が何bit続くか求めることができる。
 証明:0になるまで2倍して(=左ビットシフトして)、その回数をsから引くことで分かる。
====お気持ちパート====
この補題から「xの下位bitを見ながら順にgを構成する」という方針を思いつくが、これはうまくいかない。
例えばs=6、x=44で考える。(x=0b101100)
まずxの1bit目が0なのでgの0bit目は0。
次にxの2bit目が1なのでgの1bit目は1。
そしてxの3bit目は……と考えたいが、2bit目を0にしないと3bit目は見れない。しかし、2bit目だけを変える元(2^s-4)を見つけるのは難しい。困った……。
====お気持ちパートおわり====

そこで発想を転換して「誤差項errorを用意して、ans+ans=x+errorが常に保たれている状態のまま、errorを0にする」という方法を考える。
等式を満たす初期値としてans=error=xが取れる。
errorを0にするには下位bitから消していけば良い。
まず奇数c \in Gを勝手にとる。奇数か偶数かの判定は上述の補題により行う。Gの元の半分は奇数なので、ランダムに選ぶことで十分すばやくそのようなものが取れる。
errorのk bit目を消すためにcの2^k倍をerrorに加えるが、この時、ansにもcの2^(k-1)倍を加えることで等式ans+ans=x+errorを保ったままにすることができる。
(errorの初期値はxなので最下位bitは0でり、k>0となることに留意する)
先ほど同様s=6、x=44を例としよう。cは奇数から勝手にとって良いので例えば13とする。
c=001101 ←適当に選ぶ
error=101100 ←初期値はx
ans=101100 ←初期値はx
まずerrorの2bit目を消したいのでcを2^2倍してerror加える。同時にansにcを2^1倍して加える
c=001101
error=100000
ans=000110
次にerrorの5bit目を消したいのでcを2^5倍してerror加える。同時にansにcを2^4倍して加える
c=001101
error=000000
ans=010110
error=0となったのでこうしてans+ans=xとなるansが求められた。
(当たり前だがもう1つの解は最上位bitを反転させることで得られるのでcの2^(s-1)倍を加えれば良い)


それから位数が奇数の場合は簡単。
 補題:位数が奇数Nである巡回群Gの元g\in Gに対しては、x=\frac{N+1}{2}gとおくとx+x=gを満たす

一般の場合は、中国剰余定理を使うことで上の2つの場合に帰着できる。
 問題:位数N=q2^s (qは奇数,s>0)の巡回群(G,+)x \in Gが与えられる。g+g=xとなるg \in Gを求めよ。

先ほどと同様に等式ans+ans=x+errorを満たすようにansとerrorを更新していくことを考える。
初期値はans=\frac{q+1}{2}xerror=qxとする。
===お気持ちパート===
中国剰余定理からG\cong \mathbb{Z}/q\mathbb{Z}\times \mathbb{Z}/2^s\mathbb{Z}であるので、Gの元を0\leqq a\leqq q-10\leqq b\leqq 2^s-1により(a,b)の形で表す。
x=(x_1,x_2)とする。
ans=\frac{q+1}{2}x=(\frac{q+1}{2}x_1,\frac{q+1}{2}x_2)error=qx=(0,qx_2)である。
この時点で第1成分は既にうまく調整できているので、第2成分だけを更新していきたい。
つまりcとしては(0,奇数)となるようなものを取ってきたい。
ここで、「yがGの平方元⇔第2成分が偶数⇔N/2倍すると0」であることが容易にわかるので、cは非平方元を取ってからq倍して第1成分を潰せば良い。
===お気持ちパートおわり===
\frac{N}{2}c'\neq 0となるc'をとり、c=qc'とする。
あとは位数2ベキの場合と同様にしてcを使ってerrorの第2成分を下位bitから消していけばよい。