結論
$\sqrt{N}$以下の最大整数（整数平方根）を取得するにはmath.isqrt関数を使うとよい
経緯
先日行われたAtCoder社の競技プログラミングコンテストで以下のような出題がありました。
ABC400 C問題 2^a b^2
問題文
正の整数 $X$ は、次の条件をみたすときかつその時に限り、良い整数と呼ばれます。
- 正の整数の組 $(a,b)$ を用いて、$X=2^a \times b^2$ と書ける。
例えば、$400$ は $400=2^2 \times 10^2$ と書けるため、良い整数です。
正の整数 $N$ が与えられるので、$1$ 以上 $N$ 以下の良い整数の個数を求めてください。制約
- $1≤N≤10^{18}$
- $N$ は整数
問題文の例に上がっている$400$のように、$400=2^2 \times 10^2=2^4 \times 5^2$などと複数通りに表すことができる数があるため、愚直に数え上げると重複の管理が厄介です。
$b$を奇数に限定すればすべての良い整数を一意に表せることに着目して、以下のように実装しました。
import math
N = int(input())
ans = 0
pow = 1
n = N//(2**pow)
while n>0:
r = int(math.sqrt(n))
# 1~rまでの奇数の数
if r%2 == 0:
ans += r//2
else:
ans += (r+1)//2
pow+=1
n = N//(2**pow)
print(ans)
ところが結果はWAでした。
ただ、テストケースのほとんどで正解しており、考え方自体は悪くなさそうです。小数切り捨ての関係で不具合を出しているのかな……と考えている間に時間切れになりました。（悔しい……）
原因
公式解説がほぼ同じ方針での解法でした。本質的な差異は整数平方根を求める部分にあります。筆者がmath.sqrtを使った部分で、公式解説はmath.isqrt関数を採用していました。
そして、筆者の上述のコードの
r = int(math.sqrt(n))
を
r = math.isqrt(n)
に変更すると、いとも容易くACになりました。
このmath.isqrtとはいったい……？
Pythonのドキュメントによると、以下のような関数だそうです。
非負整数$n$の整数平方根を返します。これは$n$の正確な平方根の床であり、 $a^2 \leqq n$ を満たす最大の整数$a$と等価です。
正確な という記述が気になりますね。sqrt関数では不正確な値になる場合があるのでしょうか。簡単なコードで検証しました。
検証
$N = n^2 - 1$とします。当然$\sqrt{N} \lt n$ですから、int関数で整数部分だけを取れば$n-1$になるはずですね。しかし、実際に試すと以下のようになります。
import math
N = 10**16 - 1
print(math.sqrt(N) == 10**8) # True
print(int(math.sqrt(N))) # 100000000
print(math.isqrt(N)) # 99999999
sqrtの結果が間違っていますね。
念のため追記しますが、sqrt関数が不正確、という単純な話ではありません。実際、$N=10^{15}$くらいで試すと正しい結果になります。
$\sqrt{N}=n-(微小量)$ ですが、$N$が一定以上巨大だと微小量のオーダーが小さくなりすぎて、内部的に$0$と区別できないということでしょう。
今回の問題では$10^{18}$までのテストケースがあり得たため、見事にこのパターンに陥っていたわけです。
まとめ
正確な整数平方根を取得するには、公式ドキュメントの通り、math.sqrtではなくmath.isqrtを使うのが良いです。
なお、sqrtの誤差についてという解説がAtCoder公式から出ていて、同じ部分について書かれています。二番煎じになってしまいましたが、まあ手元で検証して学びがあったのでヨシです。