Three Sum

Solution

The brute force solution doesn't sort and takes 3 nested for loops for cubic time.

Sorting($O(n \log n)$) was not a good practice in linear time two sum but should be used in this squared time problem.

Vanilla Two Pointers

Sort and then use two pointers to search for satisfied result. Caveat is to look out for duplicates.

Complexity

time: $O(n^2)$
space: $O(n^2)$

def threeSum(self, nums: List[int]) -> List[List[int]]:
    nums.sort()
    n = len(nums)
    res = []

    for i in range(0, n-2):
        j = i+1
        k = n-1
        # avoid dup
        if i > 0 and nums[i] == nums[i-1]:
            continue
        # 2 optimizations
        if nums[i] + nums[-1] + nums[-2] < 0:
            continue
        if nums[i] + nums[i+1] + nums[i+2] > 0:
            break
        # like two-sum-ii-input-array-is-sorted
        target = -nums[i]
        while j < k:
            curr_sum = nums[j] + nums[k]
            if curr_sum == target:
                res.append([nums[i],nums[j],nums[k]])
                j += 1
                
                while j < k and nums[j] == nums[j-1]:
                    j += 1
                k -= 1
                while k > j and nums[k] == nums[k+1]:
                    k -= 1
            elif curr_sum < target:
                j += 1
            else:
                k -= 1
                
    return res

Return a Set

Also sort and then use two pointers, but the difference is returning a set rather than list.

To eliminate duplicates, use tuple rather than list for HashSet res.

Complexity

time: $O(n^2)$
space: $O(n^2)$

def threeSum(self, nums: List[int]) -> List[List[int]]:
    res = set()
    n = len(nums)
    nums.sort()

    for i in range(n-2):
        target = -nums[i]
        left = i + 1
        right = n - 1
        while left < right:
            # speed up a bit
            if nums[i] > 0:
                break
            if nums[left] + nums[right] == target:
                res.add((nums[i], nums[left], nums[right]))
                left += 1
                right -= 1
            elif nums[left] + nums[right] < target:
                left += 1
            else:
                right -= 1
    return res