AtCoder, エイシングプログラミングコンテスト 2020, C : XYZ Triplets

問題文

　https://atcoder.jp/contests/aising2020/tasks/aising2020_c

問題概要

　関数 $f$ を，$f( n )$ が以下の 2 つの条件を満たすような整数の 3 つ組 $( x, y, z )$ の個数とします．

• $1 \leq x, y, z$
• $x^2 + y^2 + z^2 + xy + yx + zx = n$

正整数 $N$ が与えられる．$f( 1 ), f( 2 ), \dots, f( N )$ をそれぞれ求めよ．

制約

• $1 \leq N$

解法

　小難しいことをする前にまずは全探索をしようということで，全探索を考えてみます．$x, y , z$ の値を全通り試す 3 重ループを考えます．このとき，各変数を $N$ まで回してしまうと $\Omega( N ^ 4 )$ になってしまって当然間に合いません．しかし，問題文の式の左辺が $N$ を超えるとどこかが break （またはループ終了）するまでずっと $N$ を超えたままです．よって，$x, y, z$ は式の左辺が $N$ 以下に収まる範囲でループシていればよいことになります．$x^2, y^2, z^2$ が入っているので，$100$ 以下の範囲でやれば十分です．また，$x, y$ が決まった時点で $x^2 + y^2 + x^y \leq N$ でなければならない……，などと考えてやると更に減らせます．このあたりまでやってコードテスト投げると最大ケースで 1.2 sec 程度になったのでこれで出しました．
　解析はちょっとむずかしいので雑にやりますが，まず，$\Theta( N )$ 回 $f$ の値を計算します．$f$ の値を求めるときは $x, y, z$ を高々 $\sqrt N$ あたりまで動かすので
\begin{align*}
O( \sqrt N ^ 3 ) &= O( ( N ^{ \frac 1 2 } )^ 3 ) \\
&= O( N ^ { 1.5 } )
\end{align*}
となります．よって，全体では $O( N ^ { 2.5 } )$ です．

もう少し賢く

　もう少しよく考えると，$f( N )$ を求めるときに $N$ 未満の各 $i$ についての $f( i )$ を一緒に求められるので，これをキャッシュしておくと指数が $1$ 減ります．

コード

int solve( const int N )
{
	int res = 0;
	for ( int i = 1; i * i <= N; ++i )
	{
		for ( int j = 1; i * i + j * j + i * j <= N; ++j )
		{
			for ( int k = 1; i * i + j * j + k * k + i * j + j * k + k * i <= N; ++k )
			{
				res += i * i + j * j + k * k + i * j + j * k + k * i == N;
			}
		}
	}
	return res;
}

int main()
{
	IN( int, N );
	REP( i, N )
	{
		cout << solve( i + 1 ) << '\n';
	}
	cout << flush;

	return 0;
}

int main()
{
	IN( int, N );

	VI res( 1'000'000 );
	REP( i, 1, 101 )
	{
		REP( j, 1, 101 )
		{
			REP( k, 1, 101 )
			{
				const int t = i * i + j * j + k * k + i * j + j * k + k * i;
				++res[t];
			}
		}
	}

	REP( i, N )
	{
		cout << res[ i + 1 ] << '\n';
	}
	cout << flush;

	return 0;
}