九章算法 | 微软面试题：三数之和II

描述

输入 n，求所有符合 x^2+y2+z^2 = n 的 x, y, z 的方案数。（x, y, z为非负整数）

• n <= 1000000
• x, y, z满足 (x<=y<=z)，只要选择出来的三个数相同就算同一种方案

在线评测地址

样例1

输入：n = 0
输出：1
解释：当其中一个为 1，剩下两个为 0，一共有 1 种方案。

样例2

输入：n = 1
输出：1
解释：当 x = 0, y = 0, z = 0 时等式成立。

找出所有的小于n的完全平方数，套3sum或者确定两个数然后套二分均可。

public class Solution {
    /**
     * @param n: The n
     * @return: The answer
     */
    public int threeSum2(int n) {
        // Write your code here
        int m = (int)Math.round(Math.sqrt(n));
        if (m * m > n) {
            m--;
        }
        int ans = 0;
        for (int i = 0; i <= m; i++) {
            int res = n - i * i;
            int left = i, right = m;
            while (left <= right) {
                int tmp = left * left + right * right;
                if (tmp > res) {
                    right--;
                } else if (tmp < res) {
                    left++;
                } else {
                    ans++;
                    left++;
                }
            }
        }
        return ans;
    }
}

更多题解参考