浮動小数点数オタクがAtCoder Beginner Contest 284のD問題をガチで解説してみる

2023/01/08に公開
1

こんにちは、浮動小数点数オタクのmod_poppoです。

昨日開催されたABC284のD問題でsqrtがどうのこうのという声がツイッターで観測されたので、ガチで考察してみます。

問題文(引用)

まず最初に問題文を引用しておきます。

問題文

正整数Nが与えられます。Nは、2つの相異なる素数p, qを用いてN=p^2qと表せることがわかっています。

p, qを求めてください。

T個のテストケースが与えられるので、それぞれについて答えを求めてください。

制約

  • 入力は全て整数
  • 1\le T\le 10
  • 1\le N\le 9\times 10^{18}
  • Nは、2つの相異なる素数p, qを用いてN=p^2qと表せる

2023が7\times 17^2と素因数分解できることにちなんだ問題のようです。

方針

まず、pqの少なくとも一方は\sqrt[3]{N}以下です(pqの両方が\sqrt[3]{N}よりも大きいと仮定すると矛盾が出ます)。制約から\sqrt[3]{N}\le \sqrt[3]{9}\times 10^6\approx 2.08\times 10^6なので、試し割りによってNの素因数を一個見つけることは難しくないでしょう。

(試し割りの際はあらかじめ素数を列挙しておいて素数だけで割るのが効率的ですが、速度を気にしないなら全ての整数で割っても良いでしょう。)

試し割りによって見つかったNの素因数をaとします。a=pなのかa=qなのかが問題です。

a=pの場合、N/aは再びaで割り切れ、q=N/a^2と計算できます。

a=qの場合、N/a=p^2なので、N/aの平方根を取ることでpが求まります。

整数除算は問題ないですが、平方根の計算は問題です。安易に浮動小数点数のsqrt関数を使ってしまうと誤差が心配です。N/aの大きさは(aは素数なのでa\ge 2を考慮して)最大で4.5\times 10^{18}\approx 2^{62}程度になり、この規模の整数は倍精度浮動小数点数で正確に表現できません。

この平方根の計算をどうするか?というのがこの記事の主題です。

平方根の計算以外の部分をC言語で実装すると次のようになります:

#include <inttypes.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

// 入力:n <= 4.5*10^18, nは平方数
int64_t isqrt(int64_t n)
{
    // どう実装する?
}

struct result {
    int64_t p, q;
};

// 入力:N <= 9*10^18
struct result solve(int64_t N)
{
    // 2080083^3 < 9*10^18 < 2080084^3
    for (int64_t a = 2; a <= 2080083; ++a) {
        if (N % a == 0) {
            int64_t b = N / a;
            if (b % a == 0) {
                // a = p
                return (struct result){.p = a, .q = b / a};
            } else {
                // a = q
                return (struct result){.p = isqrt(b), .q = a};
            }
        }
    }
    abort();
}

int main()
{
    int T;
    scanf("%d", &T);
    for (int i = 0; i < T; ++i) {
        int64_t N;
        scanf("%" SCNd64, &N);
        struct result r = solve(N);
        printf("%" PRId64 " %" PRId64 "\n", r.p, r.q);
    }
}

解法1:整数演算で済ませる

浮動小数点数の計算が心配なら、全部整数演算で片付ければ良いのです。そう、伝家の宝刀†二分探索†を使わせていただきますぞ!

// 入力:n <= 4.5*10^18, nは平方数
int64_t isqrt(int64_t n)
{
    // 4.5*10^18 < 2^62
    // 真の答えは 2 <= _ < 2^31 の範囲にある
    int64_t low = 2, high = INT64_C(1) << 31;
    while (low < high) {
        int64_t mid = (low + high) / 2;
        int64_t mid2 = mid * mid; // 最初の上界を2^31程度に抑えているのでここではオーバーフローしない
        if (mid2 < n) {
            low = mid;
        } else if (mid2 == n) {
            return mid;
        } else {
            high = mid;
        }
    }
    return low;
}

解法2:80ビットのlong doubleを使う

x86のGCCでは long double が80ビットあって、仮数部も64ビットあります。なので、64ビット整数を正確に表現できます。もちろん、四倍精度が使えるならそれでも構いません。

C言語では long double に対応する sqrtl があり、これは正確な(真の値に最も近い浮動小数点数を返す)ことが期待されるのでうまくいきます(x87にはFSQRTという命令があるのでそれを使っていることが期待できます)。

#include <assert.h>
#include <float.h>
#include <math.h>

static_assert(LDBL_MANT_DIG >= 64, "not enough precision");

// 入力:n <= 4.5*10^18, nは平方数
int64_t isqrt(int64_t n)
{
    return (int64_t)sqrtl((long double)n);
}

この方法は long double が64ビットな環境(MSVCとかApple Silicon Macとか)や long double を提供していないプログラミング言語からは使えないのがネックです。

解法3:sqrtとroundを使う

N/aが最大で2^{62}程度(なので double で正確に表現できない可能性がある)とは言っても、平方根を取ってしまえば2^{31}程度になり、この範囲の整数値は倍精度浮動小数点数で正確に表現できます。

何が言いたいかというと、sqrt(N/a) は整数pに極めて近いことが期待され、与えられた浮動小数点数に最も近い整数を返す関数 round を適用してやれば正しい答えが得られるだろう、ということです。

// 入力:n <= 4.5*10^18, nは平方数
int64_t isqrt(int64_t n)
{
    return llround(sqrt((double)n)); // llroundはroundの結果をlong longで返す関数
}

この解法はACしますが、「単にテストケースが弱い」のではなく、「本当に正しい」ことを証明してみましょう。

倍精度浮動小数点数で表現できる実数のなす集合を\mathbf{F}とし、最近接偶数丸めを行う写像をF\colon\mathbf{R}\rightarrow\mathbf{F}とします。この時、次の定理が成り立ちます:

定理1.実数xを浮動小数点数に丸める際にオーバーフロー・アンダーフローが起こらないとする。この時、次が成り立つ:

\left|\frac{F(x)-x}{x}\right|\le\frac{\varepsilon}{2}

ただし、\varepsilonは計算機イプシロン (machine epsilon) と呼ばれる値で、倍精度浮動小数点数の場合は\varepsilon=2^{-52}である。

【追記】この定理についてはより詳しい記事を書きました:浮動小数点数の丸めの相対誤差を計算機イプシロンで評価する

x>0の場合に定理1を言い換えると、(1-\varepsilon/2)x\le F(x)\le(1+\varepsilon/2)xとなります。

C言語の sqrt は大抵の処理系ではIEEE 754に準拠、つまり「真の平方根に最も近い浮動小数点数を返す」ようになっています(正確に言うならC処理系がC言語の規格のAnnex Fに準拠する場合これが成り立つ)。

なので、証明したいのはn=p^2に対して

p-\frac{1}{2}<F\bigl(\sqrt{F(n)}\bigr)<p+\frac{1}{2}

となります。

まずF(n)を定理1で評価すると

(1-\varepsilon/2)n\le F(n)\le(1+\varepsilon/2)n

となり、それぞれ平方根を取ると

\sqrt{1-\varepsilon/2}\cdot p\le\sqrt{F(n)}\le\sqrt{1+\varepsilon/2}\cdot p

となります。これを浮動小数点数に丸めると

F\bigl(\sqrt{1-\varepsilon/2}\cdot p\bigr)\le F\bigl(\sqrt{F(n)}\bigr)\le F\bigl(\sqrt{1+\varepsilon/2}\cdot p\bigr)

となり、左と右を評価すると

(1-\varepsilon/2)^{3/2}p\le F\bigl(\sqrt{F(n)}\bigr)\le (1+\varepsilon/2)^{3/2}p

となります。

あとは

p-\frac{1}{2}<(1-\varepsilon/2)^{3/2}p,\quad (1+\varepsilon/2)^{3/2}p<p+\frac{1}{2}

を示せば十分です。

前半について。実数t\ge-1について不等式(1+t)^{3/2}\geq 1+\frac{3}{2}tが成り立つことを利用すると、

\begin{aligned} (1-\varepsilon/2)^{3/2}p-\left(p-\frac{1}{2}\right) &\geq\left(1-\frac{3}{2}\cdot\frac{\varepsilon}{2}\right)p-\left(p-\frac{1}{2}\right) \\ &=\frac{1}{2}-\frac{3}{2}\cdot\frac{\varepsilon}{2} p \end{aligned}

と評価できます。あとは\varepsilon=2^{-52}p<2^{31}を使えばp-\frac{1}{2}<(1-\varepsilon/2)^{3/2}pが言えます。

後半について。実数0<t<\frac{1+\sqrt{5}}{2}について(1+t)^{3/2}<1+2tが成り立つことを使うと、

\begin{aligned} \left(p+\frac{1}{2}\right)-(1+\varepsilon/2)^{3/2}p &>p+\frac{1}{2}-(1+\varepsilon)p \\ &=\frac{1}{2}-\varepsilon p \end{aligned}

と評価でき、あとは\varepsilon pの範囲に注意すれば欲しい不等式を得ます。

よって、「倍精度浮動小数点数でsqrtしてからroundする解法」が正しいことが証明できました。

解法4:sqrtと切り捨てを使う

C系の言語では浮動小数点数を整数にキャストする際は0方向への切り捨てが行われます。解法3で round の代わりにキャストしてしまうとどうなるでしょうか?

// 入力:n <= 4.5*10^18, nは平方数
int64_t isqrt(int64_t n)
{
    return (int64_t)sqrt((double)n);
}

実はこれもACしますが、これで問題ないことを証明するのはかなり厄介です。

とりあえずこの方法が問題ないことを実験的に確認してみましょう:

#include <inttypes.h>
#include <math.h>
#include <stdio.h>

int main()
{
    // 3037000499^2 < 2^63-1 < 3037000500^2
    for (int64_t i = 0; i <= 3037000499; ++i) {
        int64_t n = i * i;
        int64_t j = (int64_t)sqrt((double)n); // sqrtの結果を切り捨てる
        if (i != j) {
            printf("%" PRId64 "\n", i);
            return 0;
        }
    }
    puts("Done");
}

このプログラムが Done を出力すれば反例はありません。(このプログラムを float / sqrtf で実行すると普通に反例が出てきます。pの大きさが浮動小数点数で正確に表現できる程度に小さいことは当たり前ですが重要です。)

では証明の時間です。

pを2乗して浮動小数点数に変換し、その平方根を取ったものを再び浮動小数点数に変換した時に元のpよりも小さくならないことを証明したいです。

実数xを浮動小数点数に変換した時に元のpより小さくならないというのは、「pの次に小さい浮動小数点数」を\mathsf{nextDown}(p)とした時に\frac{\mathsf{nextDown}(p)+p}{2}<xが成り立てば十分です。

pの指数部を整数iとして2^i< p\le 2^{i+1}とします。p=2^{i+1}の場合はp^2は浮動小数点数で正確に表現でき、その平方根も正確なので突っ込んだ議論の必要はありません。よって、評価を精密にして2^i<p\le 2^{i+1}-2^i\varepsilonの場合を考えます。

すると\mathsf{nextDown}\mathsf{nextDown}(p)=p-2^i\varepsilonと書けるので、示したいのは

p-2^{i-1}\varepsilon<\sqrt{F(p^2)}

となります。右辺は\sqrt{1-\varepsilon/2}\cdot p\le\sqrt{F(p^2)}と評価できるので、

p-2^{i-1}\varepsilon<\sqrt{1-\varepsilon/2}\cdot p

を示せば十分です。

実数0\le t<2\sqrt{2}-2\approx 0.8284について\sqrt{1-t}>1-\frac{1}{2}t-\frac{1}{4}t^2が成り立つので、

\begin{aligned} \sqrt{1-\varepsilon/2}\cdot p-\left(p-2^{i-1}\varepsilon\right) &>\left(1-\frac{1}{2}\cdot\frac{\varepsilon}{2}-\frac{1}{4}\cdot\frac{\varepsilon^2}{4}\right)p-p+2^{i-1}\varepsilon \\ &=2^{i-1}\varepsilon-\frac{\varepsilon p}{4}-\frac{\varepsilon^2 p}{16} \\ &=2^{i-1}\varepsilon\left(1-\frac{1}{2}\cdot\frac{p}{2^i}-\frac{\varepsilon p}{8\cdot 2^i}\right) \end{aligned}

ここでp\le 2^{i+1}-2^i\varepsilonを思い出すと\frac{p}{2^i}\le 2-\varepsilonなので、

\begin{aligned} \sqrt{1-\varepsilon/2}\cdot p-\left(p-2^{i-1}\varepsilon\right) &>2^{i-1}\varepsilon\left(1-\frac{1}{2}\cdot\frac{p}{2^i}-\frac{\varepsilon p}{8\cdot 2^i}\right) \\ &\ge 2^{i-1}\varepsilon\left(1-\frac{2-\varepsilon}{2}-\frac{\varepsilon p}{8\cdot 2^i}\right) \\ &= 2^{i-4}\varepsilon^2\left(4-\frac{p}{2^i}\right) \\ &>0 \end{aligned}

と評価できます。

よって、「倍精度浮動小数点数でsqrtしてから切り捨てる解法」が正しいことが証明できました。

雑感

浮動小数点数の誤差評価は面倒なのでなるべくやりたくないですね。特に切り捨ての解法の正当性の証明は面倒なので、浮動小数点数を整数に変換する際は安易なキャストではなく round 関数を使うようにしましょう。浮動小数点数を使わないやり方ができれば一番良いです。

安易にキャスト(切り捨て)して通ってしまう問題はあまり教育的ではない気もする。

過去の類似記事

Discussion

Mizar/みざーMizar/みざー

二分探索の他にも、整数上でニュートン・ラフソン法を使って平方根を計算する事もできます。

N\lt 2^{64}, 初期値s_0\lfloor\sqrt{N}\rfloor\le s_0\lt 2.9\sqrt{N} 程度の範囲であれば、 s_{n+1}=\lfloor(s_n+\lfloor a/s_n\rfloor)/2\rfloor の操作 7回目までに s_i=\lfloor\sqrt{N}\rfloor となる 0\le i \lt 7 を見つける事ができるかと思います。

https://zenn.dev/mizar/articles/791698ea860581#(参考)-整数上でのニュートン法にて平方根が計算できる理由

ログインするとコメントできます