3sum

3sum

December 14, 2023
medium

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.

Solution #

Solution 1 #

Here’s the Python code to find all unique triplets in an array that add up to zero:

def threeSum(nums):
    nums.sort()
    result = []

    for i in range(len(nums) - 2):
        # Avoid duplicate solutions
        if i > 0 and nums[i] == nums[i - 1]:
            continue

        left, right = i + 1, len(nums) - 1

        while left < right:
            total = nums[i] + nums[left] + nums[right]

            if total < 0:
                left += 1
            elif total > 0:
                right -= 1
            else:
                result.append([nums[i], nums[left], nums[right]])

                # Move left and right to the next different numbers
                while left < right and nums[left] == nums[left + 1]:
                    left += 1
                while left < right and nums[right] == nums[right - 1]:
                    right -= 1

                left += 1
                right -= 1

    return result

# Example usage
nums = [-1, 0, 1, 2, -1, -4]
three_sum_result = threeSum(nums)
print(three_sum_result)

When you run this code with the example array [-1, 0, 1, 2, -1, -4], it returns [[-1, -1, 2], [-1, 0, 1]], which are the unique triplets that sum up to zero.

Improved solution #

  1. Skip Unnecessary Iterations: If the current number is greater than 0, we can break out of the loop early because the remaining numbers will also be greater than 0, and we won’t be able to find a triplet that sums to 0.

  2. Early Stopping: If the smallest number in the sorted array is greater than 0, or the largest number is less than 0, we can return early because it’s impossible to find a triplet that sums to 0.

Here’s the improved code:

def threeSum(nums):
    nums.sort()
    result = []
    n = len(nums)

    # Early stopping if the smallest or largest number doesn't allow for a sum of 0
    if n < 3 or nums[0] > 0 or nums[-1] < 0:
        return result

    for i in range(n - 2):
        # Skip duplicate elements
        if i > 0 and nums[i] == nums[i - 1]:
            continue

        # Early break if the current number is greater than 0
        if nums[i] > 0:
            break

        left, right = i + 1, n - 1

        while left < right:
            total = nums[i] + nums[left] + nums[right]

            if total < 0:
                left += 1
            elif total > 0:
                right -= 1
            else:
                result.append([nums[i], nums[left], nums[right]])
                left += 1
                right -= 1

                # Skip duplicate elements
                while left < right and nums[left] == nums[left - 1]:
                    left += 1
                while left < right and nums[right] == nums[right + 1]:
                    right -= 1

    return result

# Example usage
nums = [-1, 0, 1, 2, -1, -4]
three_sum_result = threeSum(nums)
print(three_sum_result)

These changes help in reducing the number of unnecessary iterations, especially in cases where it’s impossible to find a valid triplet or where skipping duplicate values can save computation time.