18. 4Sum

Medium

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.

Solution

class Solution:
    def fourSum(self, nums: List[int], target: int) -> List[List[int]]:
        return ksum(4, nums, target)


def ksum(k, nums, target):
    nums.sort()

    if k == 2:
        return twoSum(nums, target)

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

    return res


# assumes nums is sorted
def twoSum(nums, target):
    n = len(nums)
    res = []

    i, j = 0, n - 1
    while i < j:
        candidate = nums[i] + nums[j]
        if candidate < target or (i > 0 and nums[i - 1] == nums[i]):
            i += 1
        elif candidate > target:
            j -= 1
        else:
            res.append([nums[i], nums[j]])
            i += 1
            j -= 1

    return res