3 Sum : Find triplets that add up to a zero

[Question]: Find triplets that add up to a zero (3 sum)

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].
Notice that the order of the output and the order of the triplets does not matter.
// TC: O(N^3)
// SC:  O(3*k)  // k is the no.of triplets
func threeSumTrippletAddUpZero(_ nums: [Int]) -> [[Int]] {
    let sorted = nums.sorted()// First sort the array
    var result = [[Int]]()
    for index in 0 ..< sorted.count - 2 {
        let first = sorted[index]
        if index > 0, sorted[index - 1] == first {
            // Iterating through the same number. Already handled.
            continue
        }

        // Now we know the first number, iterating up from the beginning, down from the end
        // to find all cases of (second + third) == - first
        var lowerIndex = index + 1
        var upperIndex = sorted.count - 1

        /// Iterate lower bound up to change the value of `second`.
        /// Increases the sum
        func nextLower() {
            let current = sorted[lowerIndex]
            repeat {
                lowerIndex += 1
            } while lowerIndex < upperIndex && sorted[lowerIndex] == current
        }

        /// Iterate upper bound down to change the value of `third`.
        /// Decreases the sum
        func nextUpper() {
            let current = sorted[upperIndex]
            repeat {
                upperIndex -= 1
            } while lowerIndex < upperIndex && sorted[upperIndex] == current
        }

        while lowerIndex < upperIndex {
            let second = sorted[lowerIndex]
            let third = sorted[upperIndex]

            switch (first + second + third) {
            case 0:
                result.append([first, second, third])
                nextLower()
                nextUpper()
            case ..<0:
                // We need to increase the sum
                nextLower()
            case 1...:
                // We need to decrease the sum
                nextUpper()
            default: fatalError()
            }
        }
    }
    return result
}

let sumTrippletInput = [-1,0,1,2,-1,-4,-2,-3,3,0,4]

let opSumTripplet = threeSumTrippletAddUpZero(sumTrippletInput)
print("Output==> ", opSumTripplet) // [[-1, 0, 1], [-1, -1, 2]]

Leave a Comment

Your email address will not be published. Required fields are marked *