【哈希表】四平方和【蓝桥杯】

题目

试题 历届真题 四平方和【第七届】【省赛】【C组】

四平方和 - 蓝桥云课

1221. 四平方和

---

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

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

输入样例：

5

输出样例：

0 0 1 2

解题

方法一：暴力 枚举(超时)

思路

从大到小暴力枚举 d、c、b、a 的所有可能值，找到第一个符合 a * a + b * b + c * c + d * d == N 的结果返回。

时间复杂度为 $O(n^4)$ 在 $5\times10^6$ 的数据规模下必定超时。

代码

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.StreamTokenizer;

public class Main {
    public static void main(String[] args) throws IOException {
        StreamTokenizer in = new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
        in.nextToken();
        int N = (int) in.nval;
        for (int d = N / 2; d > 0; --d) {
            for (int c = d; c >= 0; --c) {
                for (int b = c; b >= 0; --b) {
                    for (int a = b; a >= 0; --a) {
                        if (a * a + b * b + c * c + d * d == N) {
                            System.out.printf("%d %d %d %d\n", a, b, c, d);
                            return;
                        }
                    }
                }
            }
        }
    }
}

只枚举 d、c、b ，最后算出 a 判断是否合法。
这样做可以把时间复杂度降到 $O(n^3)$，但是还是会超时。

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.StreamTokenizer;

public class Main {
    public static void main(String[] args) throws IOException {
        StreamTokenizer in = new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
        in.nextToken();
        int N = (int) in.nval;
        for (int d = N / 2; d > 0; --d) {
            for (int c = d; c >= 0; --c) {
                for (int b = c; b >= 0; --b) {
                    int rest = N - (b * b + c * c + d * d);
                    if (rest < 0) continue;
                    int a = (int) Math.sqrt(rest);
                    if (a * a == rest) {
                        System.out.printf("%d %d %d %d\n", a, b, c, d);
                        return;
                    }
                }
            }
        }
    }
}

方法二：枚举 哈希表 二分查找 空间换时间

思路

先枚举 c、d 把合法值存入哈希表然后再枚举 a、b 找出结果可以把时间复杂度降到 $O(n^2)$。具体做法是:

维护一个哈希表(map)，键为 c 、d 的平方和，值为当 c <= d 且 c * c + d * d <= n 时 c 尽量小 d 尽量大的值，然后双层循环：外层从小到大枚举 c，内层从大到小枚举 d，为保证 c 尽量小 d 尽量大，只取第一次合法的值放入哈希表。

再开一个双层循环：外层从小到大枚举 a，内层从大到小枚举 b，算出 a 、b 的平方和，如果在哈希表中存在 n 减去 a 、b 的平方和，就说明找到了一组合法的四平方和的解，因为枚举方式保证了 a 取得尽量小 b 取得尽量大，所以第一组找到的解就是升序最大的解，从哈希表中取出 c 、d 的值直接返回。

代码

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.StreamTokenizer;
import java.util.HashMap;
import java.util.Map;

public class Main {
    public static void main(String[] args) throws IOException {
        StreamTokenizer in = new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
        in.nextToken();
        int n = (int) in.nval;
        Map<Integer, int[]> map = new HashMap<>();
        for (int c = 0; c * c * 2 <= n; ++c) {
            for (int d = c; c * c + d * d <= n; ++d) {
                map.putIfAbsent(c * c + d * d, new int[]{c, d});
            }
        }
        for (int a = 0; a * a * 4 <= n; ++a) {
            for (int b = a; a * a + b * b <= n / 2; ++b) {
                int rest = n - (a * a + b * b);
                if (map.containsKey(rest)) {
                    int[] cd = map.get(rest);
                    System.out.printf("%d %d %d %d\n", a, b, cd[0], cd[1]);
                    return;
                }
            }
        }
    }
}

二分换哈希表：

import java.util.*;
import java.io.*;

public class Main {
    public static void main(String[] args) throws IOException {
        StreamTokenizer in = new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
        in.nextToken();
        int n = (int) in.nval;
        List<int[]> arr = new ArrayList<>();
        for (int c = 0; c * c <= n; ++c) {
            for (int d = c; c * c + d * d <= n; ++d) {
                arr.add(new int[]{c * c + d * d, c, d});
            }
        }
        arr.sort((x, y) -> {
            if (x[0] != y[0]) return x[0] - y[0];
            else if (x[1] != y[1]) return x[1] - y[1];
            else return x[2] - y[2];
        });
        for (int a = 0; a * a * 4 <= n; ++a) {
            for (int b = a; (a * a + b * b) * 2 <= n; ++b) {
                int x = n - (a * a + b * b);
                int l = 0, r = arr.size() - 1;
                while (l < r) {
                    int m = l + r >> 1;
                    if (arr.get(m)[0] >= x) r = m;
                    else l = m + 1;
                }
                if (arr.get(l)[0] == x) {
                    System.out.printf("%d %d %d %d\n", a, b, arr.get(l)[1], arr.get(l)[2]);
                    return;
                }
            }
        }
    }
}