122. 四平方和（暴力枚举）
2016年蓝桥杯省赛 - 四平方和（困难）
标签：2016省赛暴力枚举
题目描述
四平方和定理，又称拉格朗日定理：每个正整数都可以表示为至多 4 个正整数的平方和。把 0 也算进去，就正好可以表示为 4 个数的平方和，比如
5=02+02+12+22 5 = 0^2 + 0^2 + 1^2 + 2^25=02+02+12+22
7=12+12+12+22 7 = 1^2 + 1^2 + 1^2 + 2^27=12+12+12+22
一个正整数可能存在多种表示法。要求把 4 个数排序为0≤a≤b≤c≤d0 \le a \le b \le c \le d0≤a≤b≤c≤d，并把所有表示法按a,b,c,da, b, c, da,b,c,d为联合主键升序排列，输出排在最前面的那一组。
输入描述
一个正整数NNN（N<5000000N < 5000000N<5000000）。
输出描述
4 个非负整数，按从小到大排序，中间用空格分开。
输入输出样例
示例 1：
输入
5输出
0 0 1 2示例 2：
输入
12输出
0 2 2 2示例 3：
输入
773535输出
1 1 267 838运行限制
- 最大运行时间：1s
- 最大运行内存: 256M
解题思路
四个数要满足a2+b2+c2+d2=Na^2 + b^2 + c^2 + d^2 = Na2+b2+c2+d2=N，并且0≤a≤b≤c≤d0 \le a \le b \le c \le d0≤a≤b≤c≤d，输出字典序最小的那一组。
直接四重循环枚举四个数太慢，但注意到：只要a,b,ca, b, ca,b,c定了，ddd也就定了——剩下的部分必须是一个完全平方数。所以先算出
d=⌊N−a2−b2−c2⌋ d = \left\lfloor \sqrt{N - a^2 - b^2 - c^2} \right\rfloord=⌊N−a2−b2−c2⌋
再回头验证a2+b2+c2+d2=Na^2 + b^2 + c^2 + d^2 = Na2+b2+c2+d2=N就行。这样四重循环直接降成三重。
把NNN看成一条长条就是：a2a^2a2、b2b^2b2、c2c^2c2之外剩下的那截，必须正好等于d2d^2d2，否则这组(a,b,c)(a,b,c)(a,b,c)作废：
枚举时每一层都从上一个数开始（b≥ab \ge ab≥a、c≥bc \ge bc≥b、d≥cd \ge cd≥c），顺序既不重复也不遗漏；同时用⋅\sqrt{\cdot}⋅收紧每层的上界，比如给定aaa后bbb最多取到N−a2\sqrt{N - a^2}N−a2，可以省掉大量不可能的取值。
又因为a,b,c,da, b, c, da,b,c,d是从小到大一层层枚举的，第一组满足等式的解，恰好就是按a,b,c,da, b, c, da,b,c,d升序排在最前面的那组，找到直接返回即可。
以样例773535走一遍：
N = 773535 a=0 整个分支试遍，无解 a=1, b=1 c=1 … 266 剩余部分都不是完全平方数 c=267 剩余 = 773535 - 1 - 1 - 267² = 702244 d = √702244 = 838，且 838 ≥ 267 → 输出 1 1 267 838代码实现
importmathdefhs(n):am=int(math.sqrt(n))# a 最大值forainrange(am+1):bm=int(math.sqrt(n-a*a))# 给定a后，b的最大值forbinrange(a,bm+1):# 保证 a <= bcm=int(math.sqrt(n-a*a-b*b))# 给定a和b后，c的最大值forcinrange(b,cm+1):# 保证 b <= ct=n-a*a-b*b-c*c# 剩下的部分应等于 d^2d=int(math.sqrt(t))# 计算difd>=canda*a+b*b+c*c+d*d==n:returna,b,c,d n=int(input())print(' '.join([str(x)forxinhs(n)]))复杂度分析
- 时间复杂度：O(N3/2)O(N^{3/2})O(N3/2)，三层枚举a,b,ca, b, ca,b,c，每层的上界都是N\sqrt{N}N级别；找到第一组解就返回，实际往往远小于这个上界。
- 空间复杂度：O(1)O(1)O(1)，只用了几个循环变量和临时变量。
