Jake Tuero

3Sum

3 min read

3Sum

Problem

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] where nums[i] + nums[j] + nums[k] == 0, and the indices i, j and k are all distinct.

The output should not contain any duplicate triplets. You may return the output and the triplets in any order.

Examples

Example 1:

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

Output: [[-1,-1,2],[-1,0,1]]

Explanation:
nums[0] + nums[1] + nums[2] = (-1) + 0 + 1 = 0.
nums[1] + nums[2] + nums[4] = 0 + 1 + (-1) = 0.
nums[0] + nums[3] + nums[4] = (-1) + 2 + (-1) = 0.
The distinct triplets are [-1,0,1] and [-1,-1,2].

Example 2:

Input: nums = [0,1,1]

Output: []

Explanation: The only possible triplet does not sum up to 0.

Example 3:

Input: nums = [0,0,0]

Output: [[0,0,0]]

Explanation: The only possible triplet sums up to 0.

Constraints

  • 3 <= nums.length <= 1000
  • -10^5 <= nums[i] <= 10^5

You should aim for a solution with O(n^2) time and O(1) space, where n is the size of the input array.

Solution

A triple sum equal of i + j + k = 0 is the same as -i = j + k. We first sort the input. We use an outer loop to to pick the element i. If this element is positive, then everything to the right is also positive and the sum will never match, so we can skip. We can also skip duplicates. We now have a subarray to search through. We set j to be right of i, and k to point to the end. If the sum is too small, we walk j forward. If the sum is too large, we walk k backward. If we have a valid sum, store, then move j and k forward/backward ensuring no duplicates.

class Solution {
public:
    vector<vector<int>> threeSum(vector<int>& nums) {
        // Sort the array
        std::sort(nums.begin(), nums.end());
        // nums[i] + nums[j] + nums[k] == 0 for i != j != k
        // -nums[i] == nums[j] + nums[k]
 
        std::vector<std::vector<int>> triplets;
        // Outer loop picks the i element
        for (std::size_t i = 0; i < nums.size(); ++i) {
            const auto el_i = nums[i];
            // If first element is positive, all other elements are positive
            // and thus we cannot get a triplet
            if (el_i > 0) {
                break;
            }
 
            // Skip duplicate starting elements
            if (i > 0 && nums[i-1] == el_i) {
                continue;
            }
 
            // j points to right of i, k points to end
            std::size_t j = i + 1; 
            std::size_t k = nums.size() - 1;
            // Walk j forward if sum too small, k backward if sum too large
            // If we get 0, then we found a triplet
            while (j < k) {
                const auto sum = el_i + nums[j] + nums[k];
                if (sum > 0) {
                    --k;
                } else if (sum < 0) {
                    ++j;
                } else {
                    // sum = 0, store
                    std::vector<int> res = {el_i, nums[j], nums[k]};
                    triplets.push_back(std::move(res));
                    // Move pointers ensuring no duplicates
                    ++j;
                    --k;
                    while (j < k && nums[j-1] == nums[j]) {
                        ++j;
                    }
                    while (j < k && nums[k+1] == nums[k]) {
                        --k;
                    }
                }
            }
        }
        return triplets;
    }
};

Graph View