2. 四平方和

题目描述

四平方和定理，又称为拉格朗日定理：
每个正整数都可以表示为至多 $4$ 个正整数的平方和。
如果把 $0$ 包括进去，就正好可以表示为 $4$ 个数的平方和。
比如：

$5=0^2+0^2+1^2+2^2$
$7=1^2+1^2+1^2+2^2$

对于一个给定的正整数，可能存在多种平方和的表示法。
要求你对 $4$ 个数排序：

$0\le a\le b\le c\le d$

并对所有的可能表示法按$a,b,c,d$为联合主键升序排列，最后输出第一个表示法。

输入格式

输入一个正整数 $N$ 。

输出格式

输出 $4$ 个非负整数，按从小到大排序，中间用空格分开。

数据范围

$0<N<5\ast 10^6$

时/空限制

$1s/64MB$

输入样例：

5

输出样例：

0 0 1 2

---

解题思路：以空间换时间来实现优化

本来，通过暴力的做法我们需要枚举 $a,b,c,d$ 四个数字，因为第四个数字可以直接通过前三个数字计算出来，所以至少需要三重循环，即 $O(n^3)$，但是这肯定是要超时的，所以就想办法优化循环的次数。

一个比较好的思路是把三重循环拆成两次二重循环。在第一次二重循环中，计算一下 $c^2+d^2$ ,然后记录下来，在第二次对 $a$ 和 $b$ 的循环中可以直接使用，而不需要再次计算。如此一来，通过空间换时间，时间复杂度就从 $O(n^3)$ 简化成了 $O(n^2)$ 。

至于如何记录 $c^2+d^2$，则可以考虑使用 <数组 + 二分> 的做法。

代码：数组 + 二分

#include<iostream>
#include<algorithm>
#include<cmath>

using namespace std;

struct nums
{
    int sum;  //sum存的是 c 和 d 的平方和
    int c;
    int d;
}A[5000000];
int idx=0;

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

int main()
{
    int n; scanf("%d",&n);
    
    for(int c=0;c<=sqrt(n);c++)
        for(int d=c;d<=sqrt(n);d++)
            A[++idx]={c*c+d*d,c,d};
            
    sort(A+1,A+1+idx,cmp);
    
    for(int a=0;a<=sqrt(n);a++)
        for(int b=a;b<=sqrt(n);b++)
        {
            int v=n-a*a-b*b;
            
            int l=1,r=idx;
            while(l<r)
            {
                int m=(l+r)/2;
                if(A[m].sum>=v)r=m;
                else l=m+1;
            }
            
            if(A[l].sum==v)
            {
                printf("%d %d %d %d",a,b,A[l].c,A[l].d);
                return 0;
            }
        }
    
    return 0;
}