0018. Four Sum

Problem

Given an array nums of n integers, return an array of all the unique quadruplets [nums[a], nums[b], nums[c], nums[d]] such that:

0 <= a, b, c, d < n
a, b, c, and d are distinct.
nums[a] + nums[b] + nums[c] + nums[d] == target

You may return the answer in any order.

Examples

Example 1:

Input: nums = [1,0,-1,0,-2,2], target = 0
Output: [[-2,-1,1,2],[-2,0,0,2],[-1,0,0,1]]

Example 2:

Input: nums = [2,2,2,2,2], target = 8
Output: [[2,2,2,2]]

Constraints:

1 <= nums.length <= 200
-109 <= nums[i] <= 109
-109 <= target <= 109

Analysis

Trs is similar to 0015.

Solution

# Solution 1: two pointer 
def solution(nums, target=0):
    nums.sort()
    res = [] 

    for h, a in enumerate(nums[:-3]):
        if h > 0 and a == nums[h - 1]: # if duplicate, pass
            continue
        # 3 sum
        nums3 = nums[h+1:]
        sum3_target = target - nums[h]
        for i, b in enumerate(nums3):
            if i > 0 and b == nums3[i - 1]: # if duplicate, pass
                continue
            # 2 sum
            j, k = i+1, len(nums3)-1
            sum2_target = sum3_target - nums3[i]
            while j < k:
                if nums3[j] + nums3[k] < sum2_target:
                    j += 1
                elif nums3[j] + nums3[k] > sum2_target:
                    k += -1
                else:
                    res.append([nums[h], nums3[i], nums3[j], nums3[k]])
                    j += 1
                    while nums3[j] == nums3[j-1] and j < k: # skip duplicates
                        j += 1
    return res

# test
nums = [1,0,-1,0,-2,2]
target = 0
print(solution(nums))

[[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]

# solution 2: k-sum
def solution(nums, target):

    def kSum(nums, target,  k):
        res = []
        
        # If we have run out of numbers to add, return res.
        if not nums:
            return res
        
        # There are k remaining values to add to the sum. The 
        # average of these values is at least target // k.
        average_value = target // k
        
        # We cannot obtain a sum of target if the smallest value
        # in nums is greater than target // k or if the largest 
        # value in nums is smaller than target // k.
        if average_value < nums[0] or nums[-1] < average_value:
            return res
        
        if k == 2:
            return twoSum(nums, target)

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

        return res

    def twoSum(nums, target):
        res = []
        l, r = 0, len(nums) - 1

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

    nums.sort()
    return kSum(nums, target, 4)

# test
nums = [1,0,-1,0,-2,2]
target = 0
print(solution(nums, target))

[[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]