四平方和c++

题目

输入样例：

输出样例：

0 0 1 2

思路

首先想到的是使用三重循环求出 a，b，c，d 可以通过 n - a - b - c 得到。理论时间复杂度为O(1000 * 1000 * 1000) = O(10^9)。因此需要想办法降低循环层数。

考虑使用两个双重循环来做。第一个双重循环保存 s = c^2+d^2 所有可能的值，又因为题目要求输出 c 和 d ，因此也要保留 c 和 d。第二个双重循环枚举所有 a 和 b 的组合，对于每个 a 和 b,判断是否有一个 s，使得 s + a^2 + b^ 2 = n。为了快速找到是否存在这样的s，可以使用二分查找。而二分查找需要序列有序，那如何对第一个双重循环求出的值进行排序呢？题目要求按升序排列，因此要在满足 s + a^2 + b^ 2 = n 的前提下，c和d尽可能小，并且c <= d;因此需要按升序进行排列。

代码

#include<bits/stdc++.h>
using namespace std;

//每个数有1 ~ sqrt(n)个可能的值，两个数可能的组合最多有（sqrt(n) * sqrt(n)）/ 2 种
const int N = 2.5e6 + 10;
struct C {
  int t, c, d;
  bool operator < (const C &s)const
  {
    if (t != s.t)return t < s.t;
    if (c != s.c)return c < s.c;
    return d < s.d;
  }
}sum[N];

int m;

int main()
{
  int n;
  scanf("%d", &n);
  
  //第一个双重循环：枚举保存c^2+d^2所有可能的值
  for (int c = 0; c * c <= n; c ++)
  for (int d = c; d * d + c * c <= n; d ++)//c <= d
  sum[m++] = {c * c + d * d, c, d};
  
  sort(sum, sum + m);
  
  //第二个双重循环
  for (int a = 0; a * a <= n; a ++)
  {
    for (int b = a; b * b + a * a <= n; b ++)
    {
      int t = n - a * a - b * b;
      int l = 0, r = m;
      while (l < r)
      {
        int mid = (l + r) >> 1;
        if(sum[mid].t >= t)r = mid;
        else l = mid + 1;
      }
      
      if (t == sum[l].t)
      {
        printf("%d %d %d %d", a, b, sum[l].c, sum[l].d);
        return 0;
      }
    }
  }
  return 0;
}