yukicoder : No.1318 ABCD quadruplets

ちょっとゴリ押し。
https://yukicoder.me/problems/no/1318

問題

整数N,Mが与えられる。
4つの整数a,b,c,dのうち、

• いずれも0以上M以下
• a^2+b^2+c^2+d^2+ab+ac+ad+bc+bd+cd が0以上N以下

となる組み合わせを求めよ。

解法

各値2乗が出てくるので、上限はO(√N)程度である。実際Mの上限も√Nに設定されている。
愚直にO(M^4)ループを回すとちょっと間に合わないので、a≦bのパターンのみ回すようにしたら間に合った。

int N,M;

int C[161616];



void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	
	for(x=0;x<=M;x++) {
		for(y=0;y<=x;y++) {
			int num=2;
			if(x==y) num=1;
			for(i=0;i<=M;i++) {
				int S=x*x+y*y+i*i+x*y+x*i+y*i;
				for(j=0;j<=M&&S+j*(x+y+j+i)<=N;j++) C[S+j*(x+y+j+i)]+=num;
			}
		}
	}
	FOR(i,N+1) cout<<C[i]<<endl;
	
}

まとめ

いや、Editorialの解法も思いついたんだけど、a≦bだけケアしたら通っちゃってびっくり。