多重循环优化

四平方和 (程序设计)

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

每个正整数都可以表示为至多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 (N<5000000)

要求输出4个非负整数，按从小到大排序，中间用空格分开

例如，输入：

5

则程序应该输出：

0 0 1 2

再例如，输入：

12

则程序应该输出：

0 2 2 2

再例如，输入：

773535

则程序应该输出：

1 1 267 838

资源约定：

峰值内存消耗（含虚拟机） < 256M

CPU消耗 < 3000ms

请严格按要求输出，不要画蛇添足地打印类似：“请您输入…” 的多余内容。

这道题，我当时就是写了一个四重循环遍历搜索，也跑了最大的一组样例，没有超时，就直接交了！后来，在网上发现，这样做好像是超时的。唉，还是自己做题做的少，算法研究的少啊！要是想在任何一门比赛中得到好成绩，还是需要算法过硬才行啊！

//1.四重循环直撸
import java.util.Scanner;

public class Main {

public static int N = 0;

public static void Op(){
for(int a=0;a*a<N;a++){
for(int b=0;b*b<N;b++){
for(int c=0;c*c<N;c++){
for(int d=0;d*d<N;d++){
if(a*a+b*b+c*c+d*d==N){
System.out.println(a+" "+b+" "+c+" "+d);
return ;
}
}
}
}
}
}

public static void main(String[] args) {
// TODO Auto-generated method stub
Scanner in = new Scanner(System.in);
while(in.hasNext()){
N = in.nextInt();
//long t1 = System.currentTimeMillis();
Op();
//long t2 = System.currentTimeMillis();
//System.out.println("Time cost:"+(t2-t1));
}
}

}

上面这个代码，虽然我试了几次都没有超时，最大的耗时是2800ms，可是毕竟是4重循环，虽然实际为O(n^2)，但是参考了书上的代码进行如下优化：

//枚举+二分搜索优化
import java.util.Arrays;
import java.util.Random;
import java.util.Scanner;

public class Main {

public static int n, line;
public static int[] arr;

public static void find(){
arr = new int[line];
for(int i=0;i<line;i++) arr[i] = i * i;
for(int a=0;a<line;a++){
for(int b=0;b<line;b++){
for(int c=0;c<line;c++){
int t = n - arr[a] - arr[b] - arr[c];
int index = -1;
index = Arrays.binarySearch(arr, t);
if(index>=0){
//  System.out.println("n:"+n);
System.out.println(a+" "+b+" "+c+" "+(int)Math.sqrt(t));
return ;
}
}

}
}
}

public static void main(String[] args) {
// TODO Auto-generated method stub
Scanner in = new Scanner(System.in);
//  Random rand = new Random();
for(int i=0;i<10000;i++){
//n = rand.nextInt(5000000);
n = in.nextInt();
long t1 = System.currentTimeMillis();
line = (int)Math.sqrt(n)+10;
find();
long t2 = System.currentTimeMillis();
//          System.out.println("Time cost:"+(t2-t1));
//          if(t2-t1>3000){
//              System.out.println("breakpoint!");
//              break;
//          }
}
}

}

优化后的算法复杂度从O（n^4)变为O(n^3*log n)，原理是这样的，既然是求 a^2 + b^2 + c^2 + d^2 = N，先枚举出[0，sqrt(N)+10]（加上10是为了保险）所有值的平方，并将其存在arr数组中(arr[i] = i * i)，然后问题就变成了“在数组arr中是否存有四个数A,B,C,D使得A+B+C+D=N”，将等式变形D=N-A-B-C，就是等价于在数组arr中寻找有没有t=N-A-B-C这样一个值，由于数组arr是升序的，可以直接使用二分查找，其复杂度为log n，再加上外围的三层循环，总复杂度为O(n^3*log n)。