15-3sum

Try it on leetcode

Description

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.

 

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

Solution(Python)

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        return self.KSum(nums, 0, 3)

    def KSum(self, nums, target, k):

        n = len(nums)

        if not n:
            return []

        average = target // k

        if average < nums[0] or average > nums[-1]:
            return []

        if k == 2:
            return self.twosum(nums, target)

        res = []

        for i in range(n):
            if i == 0 or nums[i] != nums[i - 1]:
                for subset in self.KSum(nums[i + 1:], target - nums[i], k - 1):
                    res.append([nums[i]] + subset)

        return res

    def twosum(self, nums, target):
        lo = 0
        n = len(nums)
        hi = n - 1
        res = []

        while lo < hi:
            curr_sum = nums[lo] + nums[hi]
            if curr_sum < target or (lo > 0 and nums[lo] == nums[lo - 1]):
                lo += 1
            elif curr_sum > target or (hi < len(nums) - 1 and nums[hi] == nums[hi + 1]):
                hi -= 1
            else:
                res.append([nums[lo], nums[hi]])
                lo += 1
                hi -= 1
        return res