蓝桥杯ACwing习题

https://www.acwing.com/problem/content/1223/

#include <iostream>
#include <cstring>
#include <algorithm>

using namespace std;

const int N = 2500010;

int n;

struct Sum
{
    int s , c , d;
    bool operator< (const Sum &W) const
    {
        if(s != W.s) return s < W.s;
        if(c != W.c) return c < W.c;
        return d < W.d;
    }
    
}sum[N];

int cnt;

int main()
{
    cin >> n;

    for(int a = 0 ; a * a <= n ; a ++)
       for(int b = a ; a * a + b * b <= n ; b ++)
         sum[cnt ++] = {a * a + b * b , a , b};
    
    // 空间换时间的思想 先把所有的 c d 和 c * c + d * d 的和 s 存下来
    // 时间复杂度 o(n)
    // 我们要找的是 最小的 a b c d 所以需要排序 又因为要二分 需要两端性 所以需要排序
    sort(sum , sum + cnt);
    
    // 枚举 a b  o(n) 的时间复杂度
    for(int a = 0 ; a * a <= n ; a ++)
      for(int b = a ; a * a + b * b <= n ; b ++)
      {
          int l = 0 , r = cnt - 1; // 二分存储的 c * c + d * d 的位置
          int t = n - a * a - b * b;
          while(l < r)
          {
              int mid = l + r >> 1;
              if(sum[mid].s >= t) r = mid;
              else l = mid + 1;
          }
          // r 就是第一个 >= 当前的数的位置
          if(sum[r].s == t) // 如果 == 的话 就输出  判断的原因就是 sum[mid] 有可能大 
          {
              cout << a << ' ' << b << ' ' << sum[r].c << ' ' << sum[r].d << endl;
              return 0;
          }
      }
    
    return 0;
}