How to optimize ThreeSum (LeetCode) solution?

Question

Original Problem Statement:

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.

My solution:

class Solution {
    public List<List<Integer>> threeSum(int[] nums) {
        Arrays.sort(nums); 
        List<List<Integer>> ans = new ArrayList<List<Integer>>(); 
        ArrayList<Integer> temp; 
        int left, right, sum;
        for(int i = 0; i < nums.length; i++){
            left = i+1; 
            right = nums.length - 1;
            while(left < right){
                sum = nums[i] + nums[left] + nums[right]; 
                if(sum == 0){
                    temp = new ArrayList<Integer>(); 
                    temp.add(nums[i]);
                    temp.add(nums[left]);
                    temp.add(nums[right]);
                    if(ans.contains(temp) == false){
                        ans.add(temp);
                        left++;
                        continue;
                    }
                    left++;
                    continue;
                }else if(sum<0){
                    left++; 
                    continue;
                }else{
                    right--; 
                    continue;
                }
            }
        }
        return ans; 
    }
}

I got TLE on testcase 311/312 on LeetCode, which is a really large array with big values. I've been looking for a way to optimize this, but can't seem to think of the answer. I would really appreciate any hint on what to do. Thank you.

Answer 1

There is the problem (bug) that List.contains will not be true on finding an other List instance with same values. Storing instead a Set<Integer> would have worked. However then duplicate values are removed from the set.

Your code could have been written as:

public List<List<Integer>> threeSum(int[] nums) {
    Arrays.sort(nums);
    List<int[]> answers = new ArrayList<>();
    for (int i = 0; i < nums.length; i++) {
        int left = i + 1;
        int right = nums.length - 1;
        while (left < right) {
            int sum = nums[i] + nums[left] + nums[right];
            if (sum == 0) {
                int[] temp = {nums[i], nums[left], nums[right]};
                if (!contains(answers, temp)) {
                    answers.add(temp);
                }
                left++;
            } else if (sum < 0) {
                left++;
            } else {
                right--;
            }
        }
    }
    List<List<Integer>> ans = new ArrayList<>();
    for (int[] oldAnswer : answers) {
        ans.add(List.of(oldAnswer[0], oldAnswer[1], oldAnswer[2]));
    }
    return ans;
}

private static boolean contains(List<int[]> answers, int[] answer) {
    for (int[] oldAnswer : answers) {
        if (Arrays.equals(answer, oldAnswer)) {
            return true;
        }
    }
    return false;
}

Declarations can be close to the usage, has no influence on memory usage for variables or whatever. Continues are also a vane effort.

---

For optimisations the first thing is data structures. Having a List of 3 Integers is not as good as an int[3] array. You can also use the values and hold them in variables for an entire loop.

The other optimisation is that having two values, the third sought value can be determined. In a sorted array a binary search will yield that index (or -index-1 for the insert position when not found). This binary search will tremendously speed up the algorithm as compared with moving an index in a loop toward a solution.

And then there is the short cutting loops on out-of-bound candidates. Perhaps not realized 100% but sufficiently.

public List<List<Integer>> threeSum2(int[] nums) {
    Arrays.sort(nums);
    List<int[]> answers = new ArrayList<>();
    for (int left = 0; left < nums.length - 2; left++) {
        int leftNum = nums[left];
        for (int mid = left + 1; mid < nums.length - 1; mid++) {
            int midNum = nums[mid];
            int soughtRightNum = -leftNum - midNum;
            if (soughtRightNum < midNum) {
                break;
            }
            int right = Arrays.binarySearch(nums, mid + 1, nums.length, soughtRightNum);
            if (right >= 0) { // Fpund.
                int[] ans = {left, mid, right};
                if (!contains(answers, ans)) {
                    answers.add(ans);
                }
            }
        }
    }
    List<List<Integer>> ans = new ArrayList<>();
    for (int[] oldAnswer : answers) {
        ans.add(List.of(oldAnswer[0], oldAnswer[1], oldAnswer[2]));
    }
    return ans;
}

Contains check and adding best would be done together. If there are really many answers like a thousand zeros, that would need some consideration too. That Set might be an idea.