Leetcode: 15. 3Sum

Problem Statement

Pattern

Solution

Given an array of integers \(a\), find all triplets that sums to \(0\). Here are two ways to solve this problem.

The first method uses Fix One and Search Other Template to solve this problem after sorting \(a\): for each index \(i\), find all pairs that sums to \(-a[i]\) using a variation of Leetcode: 167. Two Sum II - Input Array Is Sorted. The time complexity is \(O(n^2)\) as the two sum algorithm is \(O(n)\) and it is executed at most \(n\) times. The space complexity is \(O(\log n)\) due the sorting algorithm.

The second method does not require sorting the input but it uses extra memory to achieve the same time complexity (see Improve Performance Using More Memory). The idea is similar to Leetcode: 1. Two Sum. We use a dictionary to store indexes of values of \(a\). For each \(i\) in \(0..(n-1)\) and \(j\) in \((i+1)..(n-1)\), add a triplet to the answer if \(-a[i] - a[j]\) appears in an index different than \(i\) and \(j\). To avoid duplicates, add sorted tuples to the answer. The time complexity is \(O(n^2)\) as the solution iterates over \(n \times (n - 1) / 2\) pairs and spends \(O(1)\) to decide and add it to the solution. The space complexity is \(O(n)\) as the input values are stored in the dictionary. 

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        N = len(nums)

        seen = {}
        for i, x in enumerate(nums):
            seen.setdefault(x, set()).add(i)

        ans = set()
        for i in range(N):
            for j in range(i + 1, N):
                x = nums[i] + nums[j]
                if -x in seen and len(seen[-x]) - (i in seen[-x]) - (j in seen[-x]):
                    ans.add(tuple(sorted([nums[i], nums[j], -x])))
        return [*ans]

Cited by 1