【算法思想·数组】一个方法团灭nSum问题

Captain823Jack

已于 2024-08-31 00:13:43 修改

阅读量939

点赞数 29

文章标签：算法

于 2024-08-31 00:13:00 首次发布

本文链接：https://blog.csdn.net/Captain823Jack/article/details/141729854

版权

经常刷力扣的朋友肯定知道鼎鼎有名的 twoSum 问题，不过除了 twoSum 问题，力扣上面还有 3Sum，4Sum 问题，以后如果想出个 5Sum，6Sum 也不是不可以。

总结来说，这类 nSum 问题就是给你输入一个数组 nums 和一个目标和 target，让你从 nums 选择 n 个数，使得这些数字之和为 target。

那么，对于这种问题有没有什么好办法用套路解决呢？本文就由浅入深，层层推进，用一个函数来解决所有 nSum 类型的问题。

本文参考labuladong算法笔记[一个方法团灭 nSum 问题 | labuladong 的算法笔记]

一、twoSum 问题

我先来编一道 twoSum 题目：

如果假设输入一个数组 nums 和一个目标和 target，请你返回 nums 中能够凑出 target 的两个元素的值，比如输入 nums = [1,3,5,6], target = 9，那么算法返回两个元素 [3,6]。可以假设只有且仅有一对儿元素可以凑出 target。

我们可以先对 nums 排序，然后利用前文双指针技巧写过的左右双指针技巧，从两端相向而行就行了：

def twoSum(nums, target):
    # 先对数组排序
    nums.sort()
    # 左右指针
    l, r = 0, len(nums) - 1
    while l < r:
        sum_nums = nums[l] + nums[r]
        # 根据 sum 和 target 的比较，移动左右指针
        if sum_nums < target:
            l += 1
        elif sum_nums > target:
            r -= 1
        elif sum_nums == target:
            return [nums[l], nums[r]]
    return []

这样就可以解决这个问题，力扣第 1 题「两数之和」和力扣第 167 题「两数之和 II - 输入有序数组」稍加修改就可以用类似的思路解决，我这里就不写了。

不过我要继续魔改题目，把这个题目变得更泛化，更困难一点：

nums 中可能有多对儿元素之和都等于 target，请你的算法返回所有和为 target 的元素对儿，其中不能出现重复。

比如说输入为 nums = [1,3,1,2,2,3], target = 4，那么算法返回的结果就是：[[1,3],[2,2]]（注意，我要求返回元素，而不是索引）。

对于修改后的问题，关键难点是现在可能有多个和为 target 的数对儿，还不能重复，比如上述例子中 [1,3] 和 [3,1] 就算重复，只能算一次。

【算法思路】

首先，基本思路肯定还是排序加双指针：

def twoSum(nums, target):
    res = []
    # 先对数组排序
    nums.sort()
    # 左右指针
    l, r = 0, len(nums) - 1
    while l < r:
        sum_nums = nums[l] + nums[r]
        # 根据 sum 和 target 的比较，移动左右指针
        if sum_nums < target:
            l += 1
        elif sum_nums > target:
            r -= 1
        elif sum_nums == target:
            res.append([nums[l], nums[r])
            l -= 1
            r += 1
    return res

但是，这样实现会造成重复的结果，比如说 nums = [1,1,1,2,2,3,3], target = 4，得到的结果中 [1,3] 肯定会重复。

出问题的地方在于 sum == target 条件的 if 分支，当给 res 加入一次结果后，lo 和 hi 不仅应该相向而行，而且应该跳过所有重复的元素：

所以，可以对双指针的 while 循环做出如下修改：

def twoSumTargert(nums, target):
    # nums 数组必须有序
    nums.sort()
    l, r = 0, len(nums) - 1
    res = []
    while l < r:
        sum = nums[l] + nums[r]
        left, right = nums[l], nums[r]
        if sum < target:
            while l < r and nums[l] == left: 
                l += 1
        elif sum > target:
            while l < r and nums[r] == nums[r-1]: 
                hi -= 1
        else:
            res.append([left, right])
            while l < r and nums[l] == left: 
                l += 1
            while lo < hi and nums[r] == right: 
                r -= 1
    return res

这样，一个通用化的 twoSum 函数就写出来了，请确保你理解了该算法的逻辑，我们后面解决 3Sum 和 4Sum 的时候会复用这个函数。

这个函数的时间复杂度非常容易看出来，双指针操作的部分虽然有那么多 while 循环，但是时间复杂度还是 O(N)，而排序的时间复杂度是 O(NlogN)，所以这个函数的时间复杂度是 O(NlogN)。

二、3Sum 问题

这是力扣第 15 题「三数之和」：

15. 三数之和 | 力扣 | LeetCode |

给你一个整数数组 nums ，判断是否存在三元组 [nums[i], nums[j], nums[k]] 满足 i != j、i != k 且 j != k ，同时还满足 nums[i] + nums[j] + nums[k] == 0 。请你返回所有和为 0 且不重复的三元组。

注意：答案中不可以包含重复的三元组。

示例 1：

输入：nums = [-1,0,1,2,-1,-4]
输出：[[-1,-1,2],[-1,0,1]]
解释：
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 。
不同的三元组是 [-1,0,1] 和 [-1,-1,2] 。
注意，输出的顺序和三元组的顺序并不重要。

示例 2：

输入：nums = [0,1,1]
输出：[]
解释：唯一可能的三元组和不为 0 。

示例 3：

输入：nums = [0,0,0]
输出：[[0,0,0]]
解释：唯一可能的三元组和为 0 。

提示：

3 <= nums.length <= 3000
-105 <= nums[i] <= 105

这样，我们再泛化一下题目，不要光和为 0 的三元组了，计算和为 target 的三元组吧，同上面的 twoSum 一样，也不允许重复的结果：

【算法思路】

这个问题怎么解决呢？很简单，穷举呗。现在我们想找和为 target 的三个数字，那么对于第一个数字，可能是什么？nums 中的每一个元素 nums[i] 都有可能！

那么，确定了第一个数字之后，剩下的两个数字可以是什么呢？其实就是和为 target - nums[i] 的两个数字呗，那不就是 twoSum 函数解决的问题么？

可以直接写代码了，需要把 twoSum 函数稍作修改即可复用：

class Solution:
    # 计算数组 nums 中所有和为 target 的三元组
    def threeSumTarget(self, nums, target):
        # 数组得排个序
        nums.sort()
        n = len(nums)
        res = []
        # 穷举 threeSum 的第一个数
        for i in range(n):
            # 对 target - nums[i] 计算 twoSum
            tuples = self.twoSumTarget(nums, i + 1, target - nums[i])
            # 如果存在满足条件的二元组，再加上 nums[i] 就是结果三元组
            for tuple in tuples:
                tuple.append(nums[i])
                res.append(tuple)
            # 跳过第一个数字重复的情况，否则会出现重复结果
            while i < n - 1 and nums[i] == nums[i + 1]:
                i += 1
        return res

    # 从 nums[start] 开始，计算有序数组 nums 中所有和为 target 的二元组
    def twoSumTarget(self, nums, start, target):
        # 左指针改为从 start 开始，其他不变
        lo = start
        hi = len(nums) - 1
        res = []
        while lo < hi:
            ...
        return res

需要注意的是，类似 twoSum，3Sum 的结果也可能重复，比如输入是 nums = [1,1,1,2,3], target = 6，结果就会重复，可能选出多个 [1,2,3]。

避免重复的关键点在于，不能让第一个数重复，至于后面的两个数，我们复用的 twoSumTarget 函数会保证它们不重复。所以代码中必须用一个 while 循环来保证 3Sum 中第一个元素不重复。

至此，3Sum 问题就解决了，时间复杂度不难算，排序的复杂度为 O(NlogN)，twoSumTarget 函数中的双指针操作为 O(N)，threeSumTarget 函数在 for 循环中调用 twoSumTarget 所以总的时间复杂度就是 O(NlogN + N^2) = O(N^2)。

三、4Sum 问题

这是力扣第 18 题「四数之和」：

18. 四数之和 | 力扣 | LeetCode |

给你一个由 n 个整数组成的数组 nums ，和一个目标值 target 。请你找出并返回满足下述全部条件且不重复的四元组 [nums[a], nums[b], nums[c], nums[d]] （若两个四元组元素一一对应，则认为两个四元组重复）：

0 <= a, b, c, d < n
a、b、c 和 d 互不相同
nums[a] + nums[b] + nums[c] + nums[d] == target

你可以按 任意顺序 返回答案。

示例 1：

输入：nums = [1,0,-1,0,-2,2], target = 0
输出：[[-2,-1,1,2],[-2,0,0,2],[-1,0,0,1]]

示例 2：

输入：nums = [2,2,2,2,2], target = 8
输出：[[2,2,2,2]]

提示：

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

都到这份上了，4Sum 完全就可以用相同的思路：穷举第一个数字，然后调用 3Sum 函数计算剩下三个数，最后组合出和为 target 的四元组。

【python 解法】

class Solution:
    def fourSum(self, nums: List[int], target: int) -> List[List[int]]:
        # 数组需要排序
        nums.sort()
        n = len(nums)
        res = []
        # 穷举 fourSum 的第一个数
        for i in range(n):
            # 对 target - nums[i] 计算 threeSum
            triples = self.threeSumTarget(nums, i + 1, target - nums[i])
            # 如果存在满足条件的三元组，再加上 nums[i] 就是结果四元组
            for triple in triples:
                triple.append(nums[i])
                res.append(triple)
            # fourSum 的第一个数不能重复
            while i < n - 1 and nums[i] == nums[i + 1]: i += 1
        return res

    # 从 nums[start] 开始，计算有序数组 nums 中所有和为 target 的三元组
    def threeSumTarget(self, nums: List[int], start: int, target: int) -> List[List[int]]:
        n = len(nums)
        res = []
        # i 从 start 开始穷举，其他都不变
        for i in range(start, n):
            ...
        return res

这样，按照相同的套路，4Sum 问题就解决了，时间复杂度的分析和之前类似，for 循环中调用了 threeSumTarget 函数，所以总的时间复杂度就是 O(N^3)。

四、100Sum 问题？

在 LeetCode 上，4Sum 就到头了，但是回想刚才写 3Sum 和 4Sum 的过程，实际上是遵循相同的模式的。我相信你只要稍微修改一下 4Sum 的函数就可以复用并解决 5Sum 问题，然后解决 6Sum 问题……

那么，如果我让你求 100Sum 问题，怎么办呢？其实我们可以观察上面这些解法，统一出一个 nSum 函数：

# 注意：调用这个函数之前一定要先给 nums 排序
# n 填写想求的是几数之和，start 从哪个索引开始计算（一般填 0），target 填想凑出的目标和
def nSumTarget(nums: List[int], n: int, start: int, target: int) -> List[List[int]]:
    sz = len(nums)
    res = []
    # 至少是 2Sum，且数组大小不应该小于 n
    if n < 2 or sz < n:
        return res
    # 2Sum 是 base case
    if n == 2:
        # 双指针那一套操作
        lo, hi = start, sz-1
        while lo < hi:
            sum = nums[lo] + nums[hi]
            left, right = nums[lo], nums[hi]
            if sum < target:
                while lo < hi and nums[lo] == left:
                    lo += 1
            elif sum > target:
                while lo < hi and nums[hi] == right:
                    hi -= 1
            else:
                res.append([left, right])
                while lo < hi and nums[lo] == left:
                    lo += 1
                while lo < hi and nums[hi] == right:
                    hi -= 1
        return res
    else:
        # n > 2 时，递归计算 (n-1)Sum 的结果
        for i in range(start, sz):
            if i > start and nums[i] == nums[i - 1]:
                # 跳过重复元素
                continue
            subs = nSumTarget(nums, n-1, i+1, target-nums[i])
            for sub in subs:
                # (n-1)Sum 加上 nums[i] 就是 nSum
                sub.append(nums[i])
                res.append(sub)
        return res

嗯，看起来很长，实际上就是把之前的题目解法合并起来了，n == 2 时是 twoSum 的双指针解法，n > 2 时就是穷举第一个数字，然后递归调用计算 (n-1)Sum，组装答案。

根据之前几道题的时间复杂度可以推算，本函数的时间复杂度应该是 O(N^(n-1))，N 为数组的长度，n 为组成和的数字的个数。

需要注意的是，调用这个 nSumTarget 函数之前一定要先给 nums 数组排序，因为 nSumTarget 是一个递归函数，如果在 nSumTarget 函数里调用排序函数，那么每次递归都会进行没有必要的排序，效率会非常低。

比如说现在我们写 LeetCode 上的 4Sum 问题：

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        # n 为 3，从 nums[0] 开始计算和为 0 的三元组
        return self.nSumTarget(nums, 3, 0, 0)

    # 注意：调用这个函数之前一定要先给 nums 排序
    def nSumTarget(self, nums: List[int], n: int, start: int, target: int) -> List[List[int]]:
        # 见上文

Captain823Jack

关注

29
点赞
踩
22

收藏

觉得还不错? 一键收藏
0
评论
【算法思想·数组】一个方法团灭nSum问题

经常刷力扣的朋友肯定知道鼎鼎有名的twoSum问题，不过除了twoSum问题，力扣上面还有3Sum4Sum问题，以后如果想出个5Sum6Sum也不是不可以。总结来说，这类nSum问题就是给你输入一个数组nums和一个目标和target，让你从nums选择n个数，使得这些数字之和为target。那么，对于这种问题有没有什么好办法用套路解决呢？本文就由浅入深，层层推进，用一个函数来解决所有nSum类型的问题。
复制链接

扫一扫