Question

Given an array nums of n integers, are there elements a, b, c in nums such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note:

The solution set must not contain duplicate triplets.

Example:

Given array nums = [-1, 0, 1, 2, -1, -4],

A solution set is:
[
  [-1, 0, 1],
  [-1, -1, 2]
]

Solution

\(O(n^2)\).

from collections import Counter

class Solution:
    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        if len(nums) < 3:
            return list()

        res = list()

        neg = list()
        pos = list()
        zero = list()

        for n in nums:
            if n < 0:
                neg.append(n)
            elif n > 0:
                pos.append(n)
            else:
                zero.append(n)

        neg_set = Counter(neg)
        pos_set = Counter(pos)
        # one neg, two pos
        for n in neg_set:
            sum = -1 * n
            for m in pos_set:
                if sum - m in pos_set:
                    if m < sum - m or (m == sum - m and pos_set[m] >= 2):
                        res.append([n, m, sum - m])

        # two neg, one pos
        for m in pos_set:
            sum = -1 * m
            for n in neg_set:
                if sum - n in neg_set:
                    if n < sum - n or (n == sum - n and neg_set[n] >= 2):
                        res.append([m, n, sum - n])

        # 0, one neg, one pos
        if len(zero) >= 1:
            for n in neg_set:
                if -1 * n in pos_set:
                    res.append([0, n, -1 * n])

        if len(zero) >= 3:
            res.append([0, 0, 0])

        return res

ChatGPT Solution

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        result = []
        for i in range(len(nums) - 2):
            if i > 0 and nums[i] == nums[i - 1]:
                continue
            l, r = i + 1, len(nums) - 1
            while l < r:
                s = nums[i] + nums[l] + nums[r]
                if s == 0:
                    result.append([nums[i], nums[l], nums[r]])
                    while l < r and nums[l] == nums[l + 1]:
                        l += 1
                    while l < r and nums[r] == nums[r - 1]:
                        r -= 1
                    l += 1
                    r -= 1
                elif s < 0:
                    l += 1
                else:
                    r -= 1
        return result

This solution uses two pointers to search for the elements in the array that add up to 0. It first sorts the array, then uses a for loop to iterate through the array and fix the first element of the triplet. The two pointers l and r are then used to find the remaining two elements that add up to 0 with the fixed first element. To avoid duplicates, the solution uses two while loops to skip over duplicate elements in the array.