AtCoder Beginner Contest 420, G : sqrt(n²+n+X)

問題文

　https://atcoder.jp/contests/abc420/tasks/abc420_g

問題概要

　整数 $x$ が与えられる．$\sqrt{ n^2 + n + x }$ が整数となるような整数 $n$ を（昇順に）列挙せよ．

制約

• $|x| \leq 10^{ 14 }$

解法

　$\sqrt{ n^2 + n + x }$ が整数だとすれば，ある整数 $m $ を使って
\[
\sqrt{ n^2 + n + x } = m
\]
と書けるはずで，（根号の中が負だと訳わからんくなるので非負だとして）両辺を $2$ 乗すると
\[
n^2 + n + x = m^2
\]
になります．やや天啓的*1ですが，左辺の $n^2 + n$ を平方完成すると
\[
n^2 + n = \left( n + \frac 1 2 \right)^2 - \frac 1 4
\]
なので，元の式に代入して
\[
\left( n + \frac 1 2 \right)^2 - \frac 1 4 + x = m^2
\]
となり，明らかに邪魔な $\frac 1 4$ を払うために全体を $4$ 倍して
\[
4 \left( n + \frac 1 2 \right)^2 - 1 + 4x = 4m^2
\]
にします．指数法則によれば $(ab)^r = a^r b^r$ なので，
\begin{align*}
4 \left( n + \frac 1 2 \right)^2
&= 2^2 \left( n + \frac 1 2 \right)^2 \\
&= ( 2n + 1 )^2
\end{align*}
と変形できて，先程の式に代入して
\[
( 2n + 1 )^2 - 1 + 4x = 4m^2
\]
となります．またやや天啓的ですが色々移項すると
\[
( 2n + 1 )^2 - 4m^2 = -4x + 1
\]
となって，左辺は $2$ 乗から $2$ 乗を引いているので因数分解して
\[
( 2n + 1 + 2m )( 2n + 1 - 2m ) = -4x + 1
\]
となります．
　ここで $-4x + 1 = pq$ となるような整数 $p, q$ をとると，一般性を失わずに
\[
\left\{\begin{align*}
2n + 1 + 2m &= p \\
2n + 1 - 2m &= q
\end{align*}\right.
\]
と立式できて，上の式から下の式を引くと，
\begin{align*}
&& 2n + 1 + 2m &= p \\
-) && 2n + 1 - 2m &= q \\ \hline
&& 4m &= p - q
\end{align*}
となります．ここから $2m = \frac { p - q } 2$ が言えるので，先程の連立方程式の上の式に代入して変形すると
\begin{align*}
2n + 1 + 2m &= p \\
2n + 1 + \frac { p - q } 2 &= p \\
2n &= p - \frac { p - q } 2 - 1 \\
n &= \frac { p - \frac { p - q } 2 - 1 } 2
\end{align*}
となります．つまり，$-4x + 1$ の約数 $p$ を固定すると $-4x + 1 = pq$ となるもう一つの約数 $q$ が決まって，出力するべき整数 $\frac { p - \frac { p - q } 2 - 1 } 2$ が得られます．なので，$-4x + 1$ の（正負両方の）全ての約数に渡って上記の値を計算して重複を取り除けば，問題に正答できます．
　計算量については，$-4x + 1$ の約数の列挙に（簡単に実装すると）$\Theta( \sqrt x )$ 時間がかかり，$O( \sqrt x )$ 個ある値の重複除去に（std::sort + std::unique や std::set のような典型的な方法だと）$O( \sqrt x \log{ \sqrt x } )$ 時間がかかります．

コード

int main()
{
	IN( LL, X );

	const LL m4xp1 = -4 * X + 1;

	VLL divisors;
	for ( LL p = 1; p * p <= abs( m4xp1 ); ++p )
	{
		if ( m4xp1 % p )
		{
			continue;
		}

		const auto q = m4xp1 / p;

		divisors.PB( p );
		if ( p != q )
		{
			divisors.PB( q );
		}
	}
	divisors.reserve( SZ( divisors ) * 2 );
	transform( rbegin( divisors ), rend( divisors ), BI( divisors ), bind( multiplies< LL >(), _1, -1 ) );

	VLL res;
	FOR( p, divisors )
	{
		const auto q = m4xp1 / p;
		res.PB( ( p - ( p - q ) / 2 - 1 ) / 2 );
	}

	sort( ALL( res ) );
	res.erase( unique( ALL( res ) ), end( res ) );

	cout << SZ( res ) << endl;
	cout << res << endl;

	return 0;
}