acwing 1221. 四平方和

题目传送门

题意：对于每个整数都可以表示为四个非负整数的平方和（四平方和定理，又称为拉格朗日定理）要求给定非负整数n,输出四个非负整数a^2+b^2+c^2+d^2=n,输出字典序较小的四个数，且按从小到大排序，中间用空格分开。

思路：先确定两个数在二分搜索另外两个数。

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;
const int maxn=5e6+10;
struct node{
	int a,b;
	int sam;
};
node s[maxn];
int k=0;	
int n;
bool cmp(node a,node b)
{
	if(a.sam!=b.sam)	return a.sam<b.sam;
	if(a.a!=b.a)	return a.a<b.a;
	if(a.b!=b.b)	return a.b<b.b;
}
int search(int x)
{
	int l=0,r=k-1,mod;
	while(l<r)
	{
		mod=l+r>>1;
		if(s[mod].sam>=x)
			r=mod;
		else
			l=mod+1;
	}
	if(s[r].sam==x)
		return r;
	return -1;
}
int main()
{

	
	scanf("%d",&n);
	int m=(int)sqrt(n);
//	printf("m %d\n",m+1);
	for(int i=0;i*i<=n;i++)
		for(int j=0;j*j+i*i<=n;j++)
		{
			s[k].a=i;
			s[k].b=j;
			s[k++].sam=i*i+j*j;
		}
	sort(s,s+k,cmp);
//	for(int i=0;i<k;i++)
	//		printf("	%d %d %d\n",s[i].a,s[i].b,s[i].sam);
	for(int i=0;i*i<=n;i++)
	{
		for(int j=0;j*j+i*i<=n;j++)
		{
			int ant=search(n-i*i-j*j);	
			if(ant!=-1)
			{
				printf("%d %d %d %d\n",i,j,s[ant].a,s[ant].b);
				return 0;	
			}
		}	
	}
	return 0;
}

有一点不是太理解，最后的搜索 

	for(int i=0;i*i<=n;i++)
	{
		for(int j=0;j*j+i*i<=n;j++)
		{
			int ant=search(n-i*i-j*j);	
			if(ant!=-1)
			{
				printf("%d %d %d %d\n",i,j,s[ant].a,s[ant].b);
				return 0;	
			}
		}	
	}

换成

	for(int i=0;i<k;i++)
	{
		int ant=search(n-s[i].sam);
		if(ant!=-1)
		{
			printf("%d %d %d %d\n",s[i].a,s[i].b,s[ant].a,s[ant].b);
			return 0;	
		}
	}

应该也是正确的（），顶多会多搜几次，但是wrong 了， 不知道为什么

输出的 1 3 1228 1237也是正确的，不过字典序比0 7 167 1735要大