解説 \(1\) と同じように、 \(1 \le x + y^2, x^2 + y \le N\) を満たす正整数の組 \((x, y)\) の個数を求める問題に帰着します。

• \(x = y\) を満たす組について

\(x + y^2 = x^2 + y\) となるので、 \(1 \le x + x^2 \le N\) を満たす \(x\) の個数を数え上げればよいです。これは二分探索などによって \(O(\log N)\) 時間で求められます。

• \(x \gt y\) を満たす組について

このとき、 \(x + y^2 \le x^2 + y\) であるので、 \(1 \le x^2 + y \le N\) となる正整数の組 \((x, y)\ (x \gt y)\) を数え上げればよいです。

\(y\) を固定したとき、 \(x\) が満たすべき不等式は \(x \gt y\) かつ \(x^2 \le N - y\) 、すなわち

\[ y \lt x \le \lfloor \sqrt{N - y} \rfloor \]

となります。よって、 \(y \lt \lfloor \sqrt{N - y} \rfloor\) であるとき、求める \(x\) の個数は \(\lfloor \sqrt{N - y} \rfloor - y\) 個になります。

そして、 \(y \lt \lfloor \sqrt{N - y} \rfloor\) を満たすような \(y\) の集合は空であるか、空でなければ整数の区間 \([1, R]\ (R \ge 1)\) となります。以下では、空でないと仮定します。

求める正整数 \((x, y)\ (x \gt y)\) の個数は

\[ \begin{aligned} \sum_{y = 1}^{R} \left(\lfloor \sqrt{N - y} \rfloor - y \right) \end{aligned} \]

です。分解して、

\[ \begin{aligned} \sum_{y = 1}^{R} \left(\lfloor \sqrt{N - y} \rfloor - y \right) = \sum_{y = 1}^{R} \lfloor \sqrt{N - y} \rfloor - \sum_{y = 1}^{R} y \end{aligned} \]

とします。後者については、

\[ \begin{aligned} \sum_{y = 1}^{R} y = \frac{R (R + 1)}{2} \end{aligned} \]

で計算することができます。前者については、

\[ \begin{aligned} \sum_{y = 1}^{R} \lfloor \sqrt{N - y} \rfloor = \sum_{m = 1}^{N-1} \lfloor \sqrt{m} \rfloor - \sum_{m = 1}^{N-R-1} \lfloor \sqrt{m} \rfloor \end{aligned} \]

となるため、結局、任意の正整数 \(n\) について

\[ \begin{aligned} \sum_{m = 1}^{n} \lfloor \sqrt{m} \rfloor \end{aligned} \]

が求められればよいです。 \(k\) を \(k^2 \le n\) となる最大の非負整数とします。このとき、 \(\lfloor \sqrt{m} \rfloor = t\) となる正整数 \(m\) の個数は \((t+1)^2 - t^2 = 2t + 1\) 個あるため、

\[ \begin{aligned} \sum_{m = 1}^{n} \lfloor \sqrt{m} \rfloor &= \sum_{m = 1}^{k^2 - 1} \lfloor \sqrt{m} \rfloor + \sum_{m = k^2}^{n} \lfloor \sqrt{m} \rfloor \\ &= \sum_{t = 1}^{k-1} t (2t + 1) + \sum_{m = k^2}^{n} k \\ &= \frac{(k-1)k(4k+1)}{6} + k (n - k^2 + 1) \end{aligned} \]

となります。よって、二分探索などで \(k\) を求めることで、 \(\displaystyle \sum_{m = 1}^{n} \lfloor \sqrt{m} \rfloor\) は \(O(\log n)\) 時間で求めることができます。ただし、この値が 64bit整数型 に収まらないことがあることに注意してください。

以上から、 \(x \gt y\) の場合が計算できました。

• \(x \lt y\) の場合

対称性より、 \(x \gt y\) の場合と同様です。

以上から、テストケースあたり \(O(\log^2 N)\) または \(O(\log N)\) 時間で求めることができます。