0015. Three Sum

0015. Three Sum#

Problem#

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.

Examples#

Example 1:

Input: nums = [-1,0,1,2,-1,-4]
Output: [[-1,-1,2],[-1,0,1]]

Example 2:

Input: nums = []
Output: []

Example 3:

Input: nums = [0]
Output: []

Constraints:#

0 <= nums.length <= 3000
-105 <= nums[i] <= 105

Analysis#

We should solve this problem after we solve 0167 - Two Sum II.

  • Sort the array, which leads to \(O(nlogn)\) time complexity. This is helpful to remove duplicates because duplicates will be aligned to each other.

  • Iterate over the sorted array. Based on the current counter, we can transform the problem as a Two Sum II problem. This step in total requires \(O(n^2)\) time.

Solution#

# Solution: 
def solution(nums, target=0):
    nums.sort()
    res = [] 
    length = len(nums)

    for i, a in enumerate(nums):
        if i > 0 and a == nums[i - 1]: # if duplicate, pass
            continue

        # 2 sum
        j, k = i+1, length-1
        sum2_target = target - nums[i]
        while j < k:
            if nums[j] + nums[k] < sum2_target:
                j += 1
            elif nums[j] + nums[k] > sum2_target:
                k += -1
            else:
                res.append([nums[i], nums[j], nums[k]])
                j += 1
                while nums[j] == nums[j-1] and j < k: # skip duplicates
                    j += 1
    return res

# test
nums = [-1, 0, 1]
print(solution(nums))

nums = [-1, 0, 1, 2, -1, -4]
print(solution(nums))

nums = []
print(solution(nums))

nums = [0]
print(solution(nums))

nums = [0, 0, 0, 0]
print(solution(nums))
[[-1, 0, 1]]
[[-1, -1, 2], [-1, 0, 1]]
[]
[]
[[0, 0, 0]]