四平方和

题目

四平方和定理，又称为拉格朗日定理：每个正整数都可以表示为至多 4个正整数的平方和。
如果把 0包括进去，就正好可以表示为 4个数的平方和。
比如：
5=0^ 2 +0^2 +1 ^2+2 ^2
7=1 ^2+1 ^2+1 ^2+2 ^2
对于一个给定的正整数，可能存在多种平方和的表示法。
要求你对 4个数排序：
0≤a≤b≤c≤d
并对所有的可能表示法按 a,b,c,d为联合主键升序排列，最后输出第一个表示法。
输入格式
输入一个正整数 N。
输出格式
输出4个非负整数，按从小到大排序，中间用空格分开。
数据范围
0<N<5∗106
输入样例：
5
输出样例：
0 0 1 2

分析

题意大致明了
但是如果三重循环的话，会被卡
所以先枚举a，b，然后把a，b，a * a+b * b存起来，存到一个结构体里，然后进行排序，
然后再两重循环c，d，然后再进行二分

代码

#include<iostream>
#include<algorithm>
using namespace std;

struct node{
    int sum;
    int a;
    int b;
}aa[5678909];

bool cmp(node a,node b)
{
    if(a.sum!=b.sum) return a.sum<b.sum;
    if(a.a!=b.a) return a.a<b.a;
    return a.b<b.b;
}

int main()
{
    int n,m=0;
    cin>>n;
    for(int a=0;a*a<=n;a++)
    {
        for(int b=a;a*a+b*b<=n;b++)
        {
            aa[m++]={a*a+b*b,a,b};
        }
    }
    sort(aa,aa+m,cmp);
    for(int c=0;c*c<=n;c++)
    {
        for(int d=0;c*c+d*d<=n;d++)
        {
            int t=n-c*c-d*d;
            int l=0,r=m-1;
            while(l<r)
            {
                int mid=(l+r)/2;
                if(aa[mid].sum>=t) r=mid;
                else l=mid+1;
            }
            if(aa[l].sum==t) 
            {
                cout<<c<<" "<<d<<" "<<aa[l].a<<" "<<aa[l].b<<endl;
                return 0;
            }
        }
    }
    return 0;
}