ARC125 B - Squares

なんでコンテスト中解けなかったんだ..............

https://atcoder.jp/contests/arc125/tasks/arc125_b

問題概要

整数 N が与えられる。以下の条件を満たす整数の組 (x, y) の個数を求めよ。

  • 1 \leq x, y \leq N
  • x^2 - y は平方数

1 \leq N \leq 10^{12}

解法 

制約を見るに O(\sqrt{N}) だろうと見当がつく。

こういう問題はとりあえず条件から式を立て、式変形。

x^2 - y = k^2, k \in \mathbb{Z}_{\geq 0}

とおくと、

x^2 - y = k^2

\Leftrightarrow y = x^2 - k^2

\Leftrightarrow y = (x-k)(x+k)

1 \leq y \leq N より

1 \leq (x-k)(x+k) \leq N

ここまでは問題を見た瞬間たどり着けた。コンテスト中はこの式を睨んで 90 分以上椅子を温めることに...........（どうして......）

k \geq 0 より、 x-k \leq x+k

よって 1 \leq x-k \leq \sqrt{N}

x-k, x+k が決まれば (x, y) は一意に定まるので、 x-k を全探索してありえる x+k の個数を数えれば解けた。

ここで、 x-k と x+k の偶奇が一致していれば (x, y) は存在する。

x-k \leq x+k であることに注意して数え上げる。

反省

コンテスト中は色々迷走してしまい \Omega (N) 解法しか見えなかった。制約から明らかに O(\sqrt{N}) なんだから、それを意識しながら式を睨んだら見えたはず。範囲が絞れたのであとは全探索、は基本中の基本なのだから、そこができなかったのは本当に悔しい。

提出コード

https://atcoder.jp/contests/arc125/submissions/25292237