蓝桥杯-四平方和

四平方和定理，又称为拉格朗日定理：

每个正整数都可以表示为至多 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∗1e6

输入样例：

5

输出样例：

0 0 1 2

题解:

暴力three重循环判断会超时~

• 枚举 c 和 d 的所有可能, 并把 c, d, c * c + d * d 都存到数组中
• 对数组排序。 (从大到小, 从小到达都行, 但对应的二分不一样)
• 枚举 a 和 b 的所有可能, 用二分找 当前 a b 下, 对应的 c d, 判断是否满足 n = a * a + b * b + c * c + d * d, 满足的话输出并退出ok了

#include <bits/stdc++.h>
using namespace std;
const int N = 1e7 + 10;
struct Node {
    int c, d, sum;
};
Node v[N];
bool cmp(Node a, Node b)
{
    if (a.sum != b.sum) return a.sum < b.sum;
    if (a.c != b.c) return a.c < b.c;
    return a.d < b.d;
}

int main()
{
    int n, m = 0; cin >> n;
    for (int i = 0; i * i <= n; i ++)
        for (int j = i; i * i + j * j <= n; j ++)
            v[m ++] = {i, j, i * i + j * j};
    sort(v, v + m, cmp);  // 从小到大排序
    
    for (int i = 0; i * i <= n; i ++)
        for (int j = 0; j * j + i * i <= n; j ++)  // a b 要从 0 开始枚举, 题中要求a≤b≤c≤d
        {
            int l = 0, r = m - 1;
            int t = n - i * i - j * j;
            while (l < r)
            {
                int mid = l + r >> 1;
                if (v[mid].sum >= t) r = mid;  // 二分, 相同的v[mid].sum, 查找的是第一个满足的, 由于我们是 升序, 刚好满足 c <= d
                else l = mid + 1;
            }
            if (v[l].sum == t)
            {
                cout << i << ' ' << j << ' ' << v[l].c << ' ' << v[l].d << endl;
                return 0;
            }
        }
    return 0;
}

觉得写的不错的话, 点个赞吧~

下面不是题解, 是笔者为了让自己更好理解二分写的笔记

• 对比上面的代码可以发现 代码中有 "<--------" 注释的是发生改变的
• 其中 if (v[mid].sum <= t), 二分条件发生变化, 导致了 这个二分在 v[mid].sum 相同的情况下, 找到的是 最后一个 符合条件的 下标
• 上面的代码在 v[mid].sum 相同的情况下, 找到的是 第一个 符合条件的 下标
• so 在下面代码排序中我们是让 c 和 d 按照倒序的方法排的顺序

#include <bits/stdc++.h>
using namespace std;
const int N = 1e7 + 10;
struct Node {
    int c, d, sum;
};
Node v[N];
bool cmp(Node a, Node b)
{
    if (a.sum != b.sum) return a.sum < b.sum;
    if (a.c != b.c) return a.c > b.c;  // <------
    return a.d > b.d;       // <-------
}

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

口诀
if 后面是 r 的话 找到的是 l 的, 即 是 第一个 满足条件的
if 后面是 l 的话 找到的是 r 的, 即 是 最后一个 满足条件的