代码练习之求数组中三数和为0的所有集合

题目描述

给定一个包含 $n$ 个元素的数组 $nums$，求该数组中和为 $0$ 的所有不同三元组。
例如：
$nums=[-1,0,1,2,-1,-4]$，该数组中满足条件的三元组集合为 $[[-1,-1,2],[-1,0,1]]$。

思路

一种直观简单的思路是三重循环暴力求解，但是时间开销非常大。
我们可以将 a + b + c = 0 转化成对于每个 $a$，求解 $b$、$c$ 使得 b + c = -a。需要注意的是，题目要求所得的三元组中不含重复的，所以在代码中要跳过相同的情况，避免重复。

代码

import java.util.ArrayList;
import java.util.Arrays;

public class Solution {
    public static void main(String[] args) {
        int[] nums = {-1, 0, 1, 2, -1, -4};
        ArrayList<ArrayList<Integer>> ans = threeSum(nums);
        System.out.println(ans);
    }

    /**
     * 求出 nums 数组中所有和为0的三元组（无重复）
     * @param nums 数组
     * @return 三元组集合
     */
    private static ArrayList<ArrayList<Integer>> threeSum(int[] nums) {
        ArrayList<ArrayList<Integer>> ans = new ArrayList<>();

        int len = nums.length;
        if (len < 3)
            return ans;

        // 排序数组
        Arrays.sort(nums);

        // a + b + c = 0 ==> b + c = -a
        // 遍历 a，对于每一个 a，找到 b 和 c，使得 b + c = -a
        // a = nums[i]，b = nums[left]，c = nums[right]
        for (int i = 0; i < len - 2; i++) {
            // 跳过相同的num[i]，避免重复
            if (i != 0 && nums[i] == nums[i - 1])   continue;

            int target = -nums[i];
            // 两个滑动指针从数组两边往中间遍历
            int left = i + 1, right = len - 1;
            while (left < right) {
                // 由于 nums 数组是升序的，所以 b + c > -a 时应减少 c 的值
                if (nums[left] + nums[right] > target)
                    right--;
                else if (nums[left] + nums[right] < target)
                    left++;
                else {
                    // 找到一个三元组，加入到ans集合中
                    ArrayList<Integer> list = new ArrayList<>();
                    list.add(nums[i]);
                    list.add(nums[left]);
                    list.add(nums[right]);
                    ans.add(list);
                    // 同样的，这里跳过相同的元素也是避免重复
                    while (left + 1 < right && nums[left + 1] == nums[left])
                        left++;
                    while (right - 1 > left && nums[right - 1] == nums[right])
                        right--;
                    left++;
                    right--;
                }
            }
        }

        return ans;
    }
}