p進Newton法で平方数判定

浮動小数点数演算を使わずに高速に平方数を判定します。 (実装例はRustですがC++でも書けると思います)

```rust
use std::num::Wrapping;

pub fn is_square(x: u64) -> bool {
    let xtz = x.trailing_zeros();
    if x == 0 {
        return true;
    } else if xtz & 1 != 0 {
        return false;
    }
    let x = Wrapping(x >> xtz);

    let mut y = Wrapping(1_u64);
    for _ in 0..5 {
        y = ((Wrapping(3_u64) - y * y * x) * y) >> 1;
    }

    let mut xrt = x * y;
    if xrt.0 & (1 << 32) != 0 {
        xrt = -xrt;
    }
    let xrt = xrt.0 & ((1 << 32) - 1);

    xrt * xrt == x.0
}
```

結論から言うと `x.sqrt().round()` を2乗するほうが2倍ほど高速なようです。が、面白いので以下解説します。

## Newton法とは

Newton法は方程式 $f(x) = 0$ を解く方法のひとつです。初期値 $x_0$ を適切に定めた上で、

$$ x_{n+1} = x_n - \frac{f(x_n)}{f'(x_n)} $$

によって反復します。

Newton法は、(浮動小数点数自体の誤差を除けば)重解ではない解の近傍では1回の反復で有効桁数が2倍になることが知られています。二分法では精度は (2進表記で) 1桁ずつしか上がらないので、それに比べると圧倒的に速く収束することになります。

## Newton法による開平 (1)

たとえば、平方根を求めるには $f(x) = x^2 - c$ とおきます。すると先ほどの反復処理は

$$ x_{n+1} = \frac{1}{2}\left(x_n + \frac{c}{x_n}\right) $$

となり、比較的低コストで計算できます。実際に $\sqrt{2}$ をこの方法で反復すると、初期値を $1$ として

$$
\begin{align}
x_0 & = \mathbf{1} \\\\
x_1 & = \mathbf{1}.5 \\\\
x_2 & = \mathbf{1.41}6\ldots \\\\
x_3 & = \mathbf{1.41421}56862745098039\ldots \\\\
x_4 & = \mathbf{1.41421356237}46899106262955788901349101\ldots \\\\
x_5 & = \mathbf{1.41421356237309504880168}962350253024361\ldots \\\\
\end{align}
$$

となり、有効桁数が2倍になっていることが目視でも読み取れます。

## Newton法による開平 (2)

上の方法を実装するには除算が必要でした。除算自体もNewton法で計算できますが、実はもっとよい方法があります。平方根のかわりに**逆数平方根 (reciprocal square root)** を求める問題を考えます。つまり、 $f(x) = x^2 - c$ ではなく $f(x) = x^{-2} - c$ とおいてNewton法を適用します。すると、Newton法の反復処理は

$$ x_{n+1} = \frac{1}{2}x_n(3 - c{x_n}^2) $$

となり、除算を使わずに求めることができるようになります。最後に逆数平方根に元の数を掛けることで平方根を得られます。

## p進Newton法とは

p進Newton法はNewton法を実数上ではなくp進数上で行うものです。 (p進数のpというのは素数のことで、2進数、3進数、5進数、……という別々の具体例からなります)

### p進数

p進数の定義方法はいくつかありますが、ここでは小数展開を使った定義を紹介します。

- 実数では2進法、3進法、10進法、12進法など、同じ数を好きな表記法で表すことができますが、p進数の場合は必ずp進法で表記します。たとえば5進数は必ず5進法で表記します。
- 実数では、数の左側は必ず有限桁で止まる一方、数の右側は無限に桁があります。p進数では逆に、数の左側に無限に桁があり、数の右側は有限桁で止まります。
  - 実数(10進表記)の例: $3.48915244412\ldots$ (右側に無限に続く)
  - 5進数の例: $\ldots0240241.324$ (左側に無限に続く)
  - 実数の場合: $3$ の逆数は $0.333\ldots$ (右側に無限に続く)
  - 5進数の場合: $3$ の逆数は $\ldots1313132$ (左側に無限に続く)
- くり上がりの方向は実数と同じで、左側にくり上がります。
  - 5進数では $3 + 4 = 12$ になる。 (5余るので左にくり上がる)
- 実数の小数表記にはマイナス記号が必要ですが、p進数の(カノニカルな)表記には使いません。
  - 実数の場合: $2$ の符号を反転すると $-2$ になる。
  - 3進数の場合: $2$ の符号を反転すると $\ldots22222221$ になる。
- 実数では $0.999\ldots = 1.000\ldots$ のように複数の表記法がある場合がありますが、p進数では(左右の余分なゼロを除けば)常に同じ書き方になります。
  - 実数の場合 $03.0$ と $3$ は同じものですが、これはp進数でも同じです。
- 実数全体の集合を $\mathbb{R}$ と書くように、p進数全体の集合を $\mathbb{Q}_p$ と書きます。

一見すると不思議な定義ですが、実は実数と共通の特徴がいくつかあります。まず四則演算ができます。たとえば3進数で $2$ と $\ldots22222221$ を足すと0になることは以下のように確認できます。

```math
\begin{matrix}
& \ldots & 2 & 2 & 2 & 2 & 2 & 2 & 1 \\
 + & & & & & & & & 2 \\
\hline
& \ldots & 0 & 0 & 0 & 0 & 0 & 0 & 0
\end{matrix}
```

かけ算も同様に、通常の筆算を左に伸ばした形で考えることができます。わり算だけ難しいですが、ここでは詳細は省略します。実数と同様、0以外の数でなら必ずわり算を定義できます。

### p進絶対値・p進距離

次に重要なのが、絶対値と距離の概念があることです。実数の絶対値は非負実数で表されますが、p進数の絶対値も非負実数です。 (複素数の絶対値も非負実数ですね)

- 0の絶対値は0です。
- そうでないときは、右端から数えて何桁目に非ゼロの値があるかを調べます。
  - 小数点から右側に桁がある場合、(余分なゼロは取り除いたときの)右側の長さを $n$ として絶対値は $p^n$ です。
  - そうではない場合は、小数点から左側にゼロが $n$ 個連続しているとき絶対値を $p^{-n}$ とします。

このように定義した距離をp進絶対値といい、 $|x|\_p$ と表記します。p進絶対値は「絶対値」の要件である以下の性質を満たします。

- $x = 0 \iff |x|\_p = 0$
- $|x|\_p\cdot|y|\_p = |xy|\_p$
- $|x + y|\_p \leq \max \\{ |x|\_p, |y|\_p \\}$ (非アルキメデス性)
  - これは通常の絶対値の要件である三角不等式 ($|x + y| \leq |x| + |y|$) より強い要件です。


$|x - y|\_p$ を $x$ と $y$ の間のp進距離といいます。これは「距離」の要件を満たします。実をいうとより強く、「超距離」の条件も満たしています。これが後で誤差評価をする時に非常に有用です。

### p進整数

小数点から右側が0であるようなp進数をp進整数といいます。p進整数全体の集合を $\mathbb{Z}_p$ と書きます。 ($\mathbb{Z}/p\mathbb{Z}$ のことをこう表記する場合もありますが、本記事では別物です)

p進整数には以下の2つの特徴があります。

- 代数の観点からは、たし算、ひき算、かけ算ができます (可換環)。これは整数の特徴を引き継いでいます。
- 解析的な観点からは、 $|x|\_p \leq 1$ を満たすp進数がp進整数です。この観点からは $\mathbb{R}$ や $\mathbb{C}$ における円盤の概念と対応しています。

円盤でありながら環でもあるのがp進整数のすぐれた点です。

以下の事実をあとの証明で使います。

**事実1**. $x \in \mathbb{Q}\_p$, $n \in \mathbb{Z}$ に対して以下は同値。

- $|x|_p \leq p^{-n}$ である。
- $x = kp^n$ となる $k \in \mathbb{Z}\_p$ が存在する。

## p進整数と剰余環の関係

$\mathbb{Z}_p$ は $\mathbb{Z}/p^n\mathbb{Z}$ を無限に大きくしたもの (射影極限) とみることもできます。

**定理2**. $x \in \mathbb{Z}\_p$, $n \in \mathbb{N}$ (0も含む) とする。このとき $|x - y|\_p \leq p^{-n}$ となるような (通常の意味での) 整数 $y \in \mathbb{Z}$ が存在する。

**証明**. 小数点から左に $n$ 桁分だけを切り出せば整数になる。

**定理3**. $x \in \mathbb{Z}\_p$, $n \in \mathbb{N}$ (0も含む) とする。このとき $|x - y|\_p \leq p^{-n}$, $|x - z|\_p \leq p^{-n}$ となるような (通常の意味での) 整数 $y, z$ は $\bmod p^{n}$ で合同である。

**証明**. $p$ 進表記で右から $n$ 桁目までが等しい。これは $\bmod p^{n}$ で合同であることと同じ。

以上のことから標準的な全射 $\mathbb{Z}_p \to \mathbb{Z}/p^n\mathbb{Z}$ を考えることができます。

## p進Newton法による逆数平方根の評価

必要な反復回数を決めるために、p進Newton法での収束条件を厳密に評価していきます。以降は $p = 2$ で決め打ちです。

**定理4**. $n \in \mathbb{N}$, $c \in \mathbb{Z}\_2$, $r \in \mathbb{Q}\_2$, $cr^2 = 1$ とする (i.e. $r$ は $c$ の逆数平方根)。$x$ と $r$ の距離が $2^{-n}$ 以下であるとき、 $\frac{1}{2}x(3 - cx^2)$ と $r$ の距離は $2^{1-2n}$ 以下である。

**証明**. $k \in \mathbb{Z}\_2$ を用いて $x - r = k2^n$ と書ける。これを代入すると

```math
\begin{align}
\frac{1}{2}x(3 - cx^2) - r
& = \frac{1}{2}(r + k2^n)(3 - c(r + k2^n)^2) - r \\
& = \frac{1}{2}(r + k2^n)(3 - cr^2 - c(2r + k2^n)k2^n) - r \\
& = \frac{1}{2}(r + k2^n)(2 - c(2r + k2^n)k2^n) - r \\
& = \frac{1}{2}(2r + 2k2^n - c(r + k2^n)(2r + k2^n)k2^n) - r \\
& = \frac{1}{2}(2k2^n - c(r + k2^n)(2r + k2^n)k2^n) \\
& = (2 - c(r + k2^n)(2r + k2^n))k2^{n-1} \\
& = (2 - 2cr^2 - 3crk2^n - k^22^{2n})k2^{n-1} \\
& = -(3crk2^n + k^22^{2n})k2^{n-1} \\
& = -(3cr + k2^n)k^22^{2n-1} \\
\end{align}
```

である。ここで $(cr)^2 = c$ より ${|cr|\_2}^2 = |c|\_2 \leq 1$ であるから $|cr|\_2 \leq 1$ つまり $cr \in \mathbb{Z}\_2$ である。よって $-(3cr + k2^n)k^2$ もまた2進整数であり $|\frac{1}{2}x(3 - cx^2) - r|\_2 \leq 2^{1-2n}$ が得られる。

**定理5**. $c \in \mathbb{Z}\_2$ が逆数平方根 $r \in \mathbb{Q}\_2$ を持ち、 $|c|\_2 = 1$ を満たす ($c$ は「奇数」である) とき、$|1 - r|\_2 \leq 1/4$ または $|1 + r|\_2 \leq 1/4$ である。

**証明**. $cr^2 = 1$ より $|c|\_2{|r|\_2}^2 = 1$, ここで $|c|\_2 = 1$ より $|r|\_2 = 1$ である。よって $r \in \mathbb{Z}\_2$ がわかる。

$r$ を $\mathbb{Z}/4\mathbb{Z}$ に射影すると、$|r|\_2 = 1$ であることから $\mod 4$ では1または3のどちらかである。1ならば $|1 - r|\_2 \leq 1/4$ である。3ならば $|1 + r|\_2 \leq 1/4$ である。

以上2つをまとめると以下の評価ができます。

**定理6**. $c \in \mathbb{Z}\_2$, $|c|\_2 = 1$ が逆数平方根 $r$ を持つとする。1を初期値としてNewton法で逆数平方根を計算すると $n$ ステップ目では $2^n + 1$ 桁の精度が得られる。 (i.e. $|x_n - r|\_2 \leq 2^{-2^n - 1}$ である)

**証明**. 定理4, 5を使った帰納法。

## 整数による近似計算

$\mathbb{R}$ がコンピューターで正確に表現できないのと同様、 $\mathbb{Q}\_p$ もコンピューターで正確には表現できません。そこで、誤差を許容して近似計算をすることになります。

$\mathbb{R}$ の近似計算には浮動小数点数を使うのが一般的です。これは $x$ を $x'$ で近似したときに $|x' - x| / |x|$ が一定の値 (マシンイプシロン) に収まるというものです。また固定小数点数 ($|x' - x|$ が一定の値に収まる) を使うこともあります。

ここで絶対値としてp進絶対値を使った固定小数点数を考えてみます。つまり、 $|x' - x|\_p \leq \delta$ となるように近似します。ここで $p = 2$, $\delta = 2^{-64}$ とおくと、これは2進表記を小数点から左側64桁で打ち切る、つまり64bit整数 + 小数点の右側 ということになります。ただし小数点の右側は邪魔なので、できるだけ $\mathbb{Q}\_2$ ではなく $\mathbb{Z}\_2$ で計算するようにします。

そうすると、Newton法の誤差計算を修正する必要が出てきそうです。しかし実際にはほとんど変わりません。

**定理6'**. $c \in \mathbb{Z}\_2$, $|c|\_2 = 1$ が逆数平方根 $r$ を持つとする。1を初期値として64bit整数でNewton法で逆数平方根を計算すると $n$ ステップ目では $\min\\{2^n + 1, 63\\}$ 桁の精度が得られる。 (i.e. $|x_n - r|\_2 \leq 2^{-\min\\{2^n + 1, 63\\}}$ である)

**証明**. ほとんどの計算は $\mathbb{Z}\_2$ で行える。$\mathbb{Q}\_p$ の非アルキメデス性より、これらの計算中は桁落ち由来の誤差は $2^{-64}$ に抑えられる。最後に2で割るときに誤差が大きくなるので、各ステップの終わりでは桁落ち由来の誤差は $2^{-63}$ になる。

## 整数平方根の計算

さて、以上で2進法上の逆数平方根を求めることはできました。ここから整数上の平方根を復元する必要があります。

まず、逆数平方根に元の値 $c$ を掛けることで $\mathbb{Z}\_2$ 上の平方根(の近似)を得ることができます。元の値は正確に表現されていて、しかも $|c|\_2 = 1$ なので平方根は逆数平方根と同じ誤差で得られます。

さて、ここでもし整数平方根が存在すれば、それは $2^{-32} + 1$ 以上 $2^{32} - 1$ 以下のはずです。 ($c$ は64bitと仮定しているため) Newton法が正負どちらの平方根に収束するかは場合によるので、符号の判定は必要ですが、それを含めても下位33bitまでわかれば整数平方根の特定ができます。Newton法で得られる精度は $\min\\{2^n + 1, 63\\}$ だったので、**5回の反復**でちょうど33bitの精度が得られることになります。

実際には整数平方根がなかったり、そもそも $\mathbb{Q}\_2$ 上に平方根がない場合もありますが、その場合でもNewton法の結果として何らかの整数が得られます。わかっているのは、**「整数平方根が存在すれば、得られた値はその整数平方根である」**ということです。ですから最後に、それが本当に整数平方根であるかどうかをチェックすれば完全な判定になります。

得られた整数は十分小さいのでオーバーフローの心配も不要です。ですから、単に2乗して元の値に一致することをチェックすれば十分でした。

## 偶数の場合の処理

ここまでは全て奇数の場合を考えていました。これは $\mathbb{Q}\_2$ を扱う上で偶数のほうが扱いが難しかったからです。

偶数の場合は以下の3通りに分けられます。

- 0の場合 (この場合は常に平方数)
- $2^{2n}$ × 奇数の場合 (この場合は奇数の問題に帰着できる)
- $2^{2n+1}$ × 奇数の場合 (この場合は常に平方数ではない)

## ベンチマーク

https://github.com/qnighy/squareness に公開しました。手元 (Intel x86_64) では以下の結果になりました。

```
test bench_is_square_20         ... bench:          15 ns/iter (+/- 0)
test bench_is_square_25         ... bench:          15 ns/iter (+/- 0)
test bench_is_square_by_sqrt_20 ... bench:           9 ns/iter (+/- 1)
test bench_is_square_by_sqrt_25 ... bench:           7 ns/iter (+/- 1)
```

`is_square` が今回説明した実装、 `is_square_by_sqrt` は浮動小数点数の `sqrt` を使った実装 (平方根を取って整数に丸め、二乗して元に戻ったらtrue) です。64bit整数の整数平方根を当てるには32bit程度の精度があれば十分なので52bitあれば十分なはず。

とにかく、現状の実装では浮動小数点数を使った実装に軍配が上がりました。

## まとめ

- Newton法は方程式を解くのに使える強力な反復解法
- p進Newton法はNewton法を実数ではなくp進数で行う
- p進数は差の大きさではなく差の割り切れ具合を距離として収束を定義する数論みの強い数
- p進数の距離は超距離になるので誤差評価が実数より楽
- 2進整数は計算機との相性がよい