[Leetcode] 3Sum

Given an array S of n integers, are there elements a, b, c in S such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note:

• Elements in a triplet (a,b,c) must be in non-descending order. (ie, a ≤ b ≤ c)
• The solution set must not contain duplicate triplets.

    For example, given array S = {-1 0 1 2 -1 -4},

    A solution set is:
    (-1, 0, 1)
    (-1, -1, 2)

We first sort the array, then enumerate the first element of the three integers, then we look at the right part of the first element, that is the subarray with the element greater or equal with the element. We start from the left-most part and right-most part of this subarray, slowly narrow this "window" to try to find the solution.

1. if num[i] + num[left] + num[right] == 0, we find a solution
2. if num[i] + num[left] + num[right] < 0, we know that we could not move the right part, since it would make the sum even smaller, so we need to move left
3. if num[i] + num[left] + num[right] > 0, we could not move left so we move right

In the process, when the new swept element is the same with the prior one, we should skip this element otherwise it would generate duplicate triplets.

All other K-sum problem could be solve in this idea, first sort the array, then enumerate all leading k - 2 elements, then find the remaining two elements using this narrowing window method. The final time complexity is O(N^(k-1)).

public class Solution {
    public ArrayList<ArrayList<Integer>> threeSum(int[] num) {
        ArrayList<ArrayList<Integer>> res = new ArrayList<ArrayList<Integer>>();
        int n = num.length;
        if(n < 3)
            return res;
        
        Arrays.sort(num);
        int pre = num[0];
        for(int i = 0; i < n-2; i++) {
            if(i > 0 && pre == num[i])
                continue;
            pre = num[i];
            int left = i + 1;
            int right = n - 1;
            while(left < right) {
                int sum = num[i] + num[left] + num[right];
                if(sum == 0) {
                    ArrayList<Integer> list = new ArrayList<Integer>();
                    list.add(num[i]);
                    list.add(num[left]);
                    list.add(num[right]);
                    res.add(list);
                    left++;
                    while(left < right && num[left] == num[left-1])
                        left++;
                }
                else if(sum > 0) {
                    right--;
                    while(left < right && num[right] == num[right+1])
                        right--;
                }
                else {
                    left++;
                    while(left < right && num[left] == num[left-1])
                        left++;
                }
            }
        }
        
        return res;
    }
}