15. 3Sum

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.

Solution

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        ret = []
        idx = 0
        while idx < len(nums)-2:
            while idx !=0 and idx<len(nums) and nums[idx-1]==nums[idx]:
                idx+=1
            if idx < len(nums):    
                target = -1 * nums[idx]
                l, r = idx+1, len(nums)-1
                while l < r:
                    if nums[l] + nums[r] == target:
                        ret.append([nums[idx], nums[l], nums[r]])
                        l+=1
                        while l < r and nums[l-1]==nums[l]:
                            l+=1
                    elif nums[l] + nums[r] < target:
                        l+=1
                    else:
                        r-=1
                idx+=1
        return ret

Time Complexity: O(n^2)

Space Complexity: O(1)

Explanation

The idea is to first sort the array and then linearly traverse the, now sorted, array and 2sum the piece of the array to the right of the current value.

The key to avoid duplicates is to shift the first value until it is different than the last. Additionally, when processing the piece of the array (2sum part) shift until a new value is found as well (only need to when you find a valid triplet).