isqrtのアルゴリズムについて

Question

isqrtのアルゴリズムの実装について

python math.isqrt
pythonのmath.isqrtは正の数nに対して、a**2がnを超えない最大の整数aを関数です。

この関数の実装が気になったので、bisectを使ったセルフ実装とpythonの実装を比べて遊んでいました。(当然ですが)math.isqrtのほうが速かったので、ソースを眺めてみました。
cpython/Modules/mathmodule.c (1555行目~)
等価な実装として、以下のurlも貼られています。
mdickinson/snippets/proofs/isqrt/src/isqrt.lean

  def isqrt_aux(c, n):
      if c == 0:
          return 1
      else:
          k = (c - 1) // 2
          a = isqrt_aux(c // 2, n >> 2*k + 2)
          return (a << k) + (n >> k+2) // a # (A)

  def isqrt(n):
      if n == 0:
          return 0
      else:
          a = isqrt_aux((n.bit_length() - 1) // 2, n)
          return a - 1 if n < a * a else a # (B)

リンク先に証明は記述されていますが、なかなか理解には至りませんでした。
大枠としてつかんだ理解では、ビット長cの数nに対する問題を、nを(大体)ビット長の半分だけ右シフト(4^kオーダーの数を2^kオーダーに)した数に関する問題として再帰的に解くことができるということです。

ただ、コードの(A)の第一項、第二項が具体的に何を表している数なのか、(B)のチェックがなぜ必要なのかを理解できていません。

アルゴリズムに詳しい方、よろしくお願いします。

Answer 1

• ニュートン法と関数のスケーリングで考える

wikepedia ニュートン法

関数 y = x^2 - C ... (*) を考える。
ニュートン法から、逐次的に漸近解を求めるには、
x_n+1 = x_n - f(x_n)/f'(x_n)
今考えている二次関数について、
x_n+1 = 1/2 * (x_n + C/x_n) ... (**)

ここで、小さな値についての解を計算し、それをスケールさせてもとの解を求めることを考える。
たとえば、Cを1/N倍したC/Nを例にとる。
(*)の式の辺々をNで割る。
y/N = y_ = x^2 / N + C/N
= (x/sqrt(N))^2 + C/N
= x_^2 + C/N

Cを1/Nだけスケールさせたとき、xは1/sqrt(N)だけスケールすることになる。
スケールした式でx_n+1が計算できるなら、その値をsqrt(N)倍すれば欲しい解が求まるはずである。

sqrtの計算に都合のいいように、Cを1/2^(2k + 2)だけスケールさせる(これは2K+2だけ右シフトすることと同じ)。Xは2^-(k + 1)だけスケールする。スケールした先で求まった漸近解をaとすると、(**)は、
x_n+1 = 1/2 * (a + (C >> 2*k + 2)/a)
x_n+1をもとのスケールに戻すには、2^(k+1)だけもどしてやればいいので、
1/2の係数を考慮して一項目のaはkだけ左シフト、二項目はk + 2だけ左シフトすればいい。

最後に、Cが平方数であるかどうかをチェックして、そうでなければ-1した値が条件を満たす最大のということになる。