15. 3Sum

Published on September 11, 2023

Question

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.

Notice that the solution set must not contain duplicate triplets.

Algorithm

An follow up of 2 sum. You have three numbers and let them sum to 0. This can be transformed to 2 sum question. For each number in the nums array, you need to find if the if there are other two numbers that sum to 0 - num.

Then this becomes a seen question.

One thing to notice is the to avoid repeat, you need to sort the array so to skip repeat numbers.

Code

class Solution {
    Set<List<Integer>> res = new HashSet<>();
    public List<List<Integer>> threeSum(int[] nums) {
        Arrays.sort(nums);
        for (int i = 0; i < nums.length; i++) {
            if (i > 0 && nums[i] == nums[i - 1]) continue;
            helper(nums, i);
        }
        return new ArrayList<>(res);
    }
    
    private void helper(int[] nums, int index) {
        int target = -nums[index];
        Set<Integer> set = new HashSet<>();
        // can be binary search since nums is sorted;
        for (int i = index + 1; i < nums.length; i++) {
            if (set.contains(target - nums[i])) {
                List<Integer> list = new ArrayList<>();
                list.add(nums[index]);
                list.add(nums[i]);
                list.add(target - nums[i]);
                res.add(new ArrayList<>(list));
            } else {
                set.add(nums[i]);
            }
        }
    }
    
}

This code can be optimized, since we are using sorted nums, we can utilize the binary search in the nums to find the other number in the helper function.

class Solution {
    public List<List<Integer>> threeSum(int[] nums) {
        Arrays.sort(nums);
        List<List<Integer>> res = new ArrayList<>();
        for (int i = 0; i < nums.length; i++) {
            if (i == 0 || nums[i - 1] != nums[i]) {
                twoSum(res, nums, i);
            }
        }
        return res;
    }
    
    private void twoSum(List<List<Integer>> res, int[] nums, int i) {
        Set<Integer> set = new HashSet<>();
        int lo = i + 1, hi = nums.length - 1;
        int sum = -nums[i];
        while (lo < hi) {
            if (nums[lo] + nums[hi] == sum) {
                res.add(Arrays.asList(nums[lo], nums[hi], nums[i]));
                while (lo + 1 < hi && nums[lo] == nums[lo + 1]) lo++;
                while (hi - 1 >= 0 && nums[hi] == nums[hi - 1]) hi--;
                lo++;
                hi--;
            } else if (nums[lo] + nums[hi] > sum) {
                hi--;
            } else {
                lo++;
            }
        }
    }
}