信息奥赛题解｜四平方和

🚀 题目浏览

【题目描述】

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

每个正整数都可以表示为至多 $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 \lt N \lt 5∗10^6$

【输入样例】

【输出样例】

0 0 1 2

【原题链接】

https://www.luogu.com.cn/problem/P8635

☘️ 题解分析

四重循环的暴力枚举做法，显然会 TLE，所以可以采用哈希的方法，来降低时间复杂度。

正确思路：

将 $c$ 和 $d$ 的平方和存储到自己模拟的哈希表中，该步复杂度为 $O(\sqrt n) * O(\sqrt n) = O(n)$
枚举 $a，b$，并且在哈希表中查找 $n - a * a - b * b$，该步复杂度为 $(O\sqrt n) * O(\sqrt n) * O(1) = O(n)$

所以该思路的时间复杂度为 $O(n) + O(n) = O(n)$，满足该题的数据范围。

本题推荐使用自己 用数组模拟的哈希表（相较于 STL 会更加快）

🧑🏻‍💻 C++ 代码

#include<bits/stdc++.h>

using namespace std;
const int N = 5e6 + 10;

int n;
int C[N], D[N];  // 哈希表，C[k]存储平方和为k时，c的值；D[k]存储平方和为k时，d的值

int main() {
    cin >> n;

    // 将c、d的平方和存入哈希表（复杂度为O(N))）
    memset(C, -1, sizeof(C)); // 初始化为-1，因为0是有实际含义的
    memset(D, -1, sizeof(D));
    for (int c = 0; c * c <= n; c++) {
        for (int d = c; c * c + d * d <= n; d++) {
            int sum = c * c + d * d;
            if (C[sum] == -1)  // 该总和第一次出现,记录此时c和d的值
                C[sum] = c, D[sum] = d;
        }
    }

    // 枚举a，b，查找 n - a*a - b*b 的哈希值
    // 哈希值存在，说明此时a,b,c,d平方和为n
    // 复杂度是sqrt(n) * sqrt(n) * O(1)= O(n) 哈希表查找为O(1)
    for (int a = 0; a * a <= n; a++) {
        for (int b = a; a * a + b * b <= n; b++) {
            int dis = n - a * a - b * b;
            if (C[dis] > -1) {
                cout << a << " " << b << " " << C[dis] << " " << D[dis] << endl;
                return 0;  // 下面没有更多需求的话，直接return 0结束即可，不用写goto
            }
        }
    }

    return 0;
}