条件を満たす正整数の組 \((a, b)\) に対して、 \(x + y^2 = a\) かつ \(x^2 + y = b\) を満たす正整数の組 \((x, y)\) はただ \(1\) つだけであることが分かります。

したがって、この問題は次のように言い換えられます： \(1 \le x + y^2, x^2 + y \le N\) を満たす正整数の組 \((x, y)\) の個数を求めてください。

この問題に帰着することができれば、あとは様々な方法でこの問題を解くことができます。

解法

\(k\) を、 \(k^2 + k \le N\) を満たす最大の非負整数とします。すると、答えは \(k^2 + 2 \times \max\{0, N - (k+1)^2\}\) になります。よって、 \(k\) を二分探索などで求めることで、テストケースあたり \(O(\log N)\) で解くことができます。

以下では、この証明を行います。実際に問題を解くときは、答えを出力する愚直なコードを書き、 \(N = 1, 2, 3, \cdots\) の順に出力して観察すると、答えの式を予想することが可能です。

証明

「条件」と言うと、 \(1 \le x + y^2, x^2 + y \le N\) を指すこととします。

• \(1 \le x, y \le k\) を満たす正整数の組 \((x, y)\) はすべて条件を満たし、合計で \(k^2\) 個ある

\(1 \le x + y^2, x^2 + y \le k^2 + k \le N\) より従います。

• \(x \ge k+2\) または \(y \ge k+2\) を満たす任意の正整数の組 \((x, y)\) については条件を満たさない

条件式の対称性から、一般性を失わず \(x \ge y\) とできます。そうすると、

\[ \begin{aligned} x - y &\le (x - y)(x + y) \\ &= x^2 - y^2 \end{aligned} \]

より \(x + y^2 \le x^2 + y\) であり、よって

\[ \begin{aligned} x^2 + y &\ge (k+2)^2\\ &= (k+1)^2 + 2(k+1) + 1\\ &\ge (k+1)^2 + (k+1)\\ &\gt N \end{aligned} \]

を満たします。最後の不等号では、「 \(k\) は \(k^2 + k \le N\) を満たす最大の非負整数である」という仮定より \(N \lt (k+1)^2 + (k+1)\) であることを使っています。 \(x^2 + y \gt N\) より、 \((x, y)\) は条件を満たしません。

• \(x = k+1\) または \(y = k+1\) を満たす正整数の組 \((x, y)\) は \(2 \times \max\{0, N - (k+1)^2\}\) 個ある

条件を満たし、 \(x = k+1\) または \(y = k+1\) であるような正整数の組 \((x, y)\) の個数を求めます。

条件式の対称性から、一般性を失わず \(x \ge y\) とできます。そうすると、 \(x = k+1\) かつ \(y \le k+1\) の場合のみを考えればよいです。いま、 \(x \ge y\) なので \(x + y^2 \le x^2 + y\) です。よって、 \(x^2 + y \le N\) か否かだけが重要になります。

\(x = k+1\) なので、

\[ \begin{aligned} y &\le N - x^2 \\ &= N - (k+1)^2 \end{aligned} \]

となります。 \(k\) の仮定から \(N \lt (k+1)^2 + (k+1)\) であるので、 \(N - (k+1)^2 \lt k+1\) です。よって、条件を満たす正整数 \(y\) の個数は \(\max\{0, N - (k+1)^2\}\) 個です。

上で自動的に \(x = k+1 > y\) が分かったので、 \(x = y\) となる場合は考えなくてよいです。以上から、 \(x = k+1\) または \(y = k+1\) であるような正整数の組 \((x, y)\) の個数は \(2 \times \max\{0, N - (k+1)^2\}\) 個になります。