经常刷力扣的朋友肯定知道鼎鼎有名的 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 题「三数之和」:
给你一个整数数组 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 题「四数之和」:
给你一个由 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]]:
# 见上文