K数之和(K-sum)问题
本文中题目出处为LeetCode 1. 两数之和, LeetCode 15. 三数之和 和 LeetCode 18. 四数之和
K数之和问题可以用一句话描述:在N个数中找出K个数之和等于输入值。该问题的子集在LeetCode中出现多次,所以在这里一个总结。
-
两数之和
- 问题
给定一个整数数组 nums 和一个目标值 target,请你在该数组中找出和为目标值的那两个整数,并返回他们的数组下标。
你可以假设每种输入只会对应一个答案。但是,你不能重复利用这个数组中同样的元素。
- 示例
给定 nums = [2, 7, 11, 15], target = 9 因为 nums[0] + nums[1] = 2 + 7 = 9 所以返回 [0, 1]
``
- 解法
本题为简单难度,如果用暴力枚举法时间复杂度为O(N^2)。
可以通过一次扫描整个列表求解,使用哈希表记录出现过的数字和对应的位置,时间复杂度为O(N)。
class Solution: def twoSum(self, nums: List[int], target: int) -> List[int]: # dict value -> index tmp_dict = {} for i, n in enumerate(nums): if target - n in tmp_dict: # return [index0, index1] return [tmp_dict[target - n], i] else: # store value -> index tmp_dict[n] = i return []
``
-
三数之和
- 问题
给定一个包含 n 个整数的数组 nums,判断 nums 中是否存在三个元素 a,b,c ,使得 a + b + c = 0 ?找出所有满足条件且不重复的三元组。
注意:答案中不可以包含重复的三元组。
- 示例
给定数组 nums = [-1, 0, 1, 2, -1, -4], 满足要求的三元组集合为: [ [-1, 0, 1], [-1, -1, 2] ]
``
- 解法
本题为中等难度,主要思想是先对列表进行排序,然后从左至右扫描列表中每一个数字,设置两个指针,对于每个数字,左指针指向其右侧第一个数字,右指针指向列表末尾。每次循环求这三个数字的和,如果和小于0,则左指针右移一位,如果和大于0,则右指针左移一位,如果和等于0,则把三个数字加入到结果中,除此以外还要跳过重复的数字。时间复杂度O(N^2),如果用暴力解法则需要O(N^3)复杂度。
class Solution: def threeSum(self, nums: List[int]) -> List[List[int]]: result = [] nums.sort() i = 0 while i < len(nums) - 2: if nums[i] > 0: break # 左指针 left = i + 1 # 右指针 right = len(nums) - 1 while left < right: tmp = [nums[i], nums[left], nums[right]] _sum = sum(tmp) if _sum == 0: result.append(tmp) right -= 1 left += 1 # 跳过左侧重复的数字 while left < right and nums[left] == nums[left - 1]: left += 1 # 跳过右侧重复的数字 while left < right and nums[right] == nums[right + 1]: right -= 1 elif _sum > 0: right -= 1 else: left += 1 # 跳过重复元素 i += 1 while i < len(nums) - 2 and nums[i - 1] == nums[i]: i += 1 return result
``
-
四数之和
- 问题
给定一个包含 n 个整数的数组 nums 和一个目标值 target,判断 nums 中是否存在四个元素 a,b,c 和 d ,使得 a + b + c + d 的值与 target 相等?找出所有满足条件且不重复的四元组。
注意:
答案中不可以包含重复的四元组。
- 示例
给定数组 nums = [1, 0, -1, 0, -2, 2],和 target = 0。 满足要求的四元组集合为: [ [-1, 0, 0, 1], [-2, -1, 1, 2], [-2, 0, 0, 2] ]
``
- 解法
本题为中等难度,解出了三数之和后很容易就可以将解法拓展到四数之和,甚至M数之和。四数之和也是先排序,然后从左至右依次扫描列表中每一个数字,假设某个数字为X,则其右侧列表中的数字需要满足三数之和等于(target - X)。所以四数之和的每次循环是三数之和问题,总复杂度为O(N^3)。
拓展到K数之和,当K>3时,可以将问题分解为K-1数之和,K-2数之和,直到三数之和,所以M数之和的复杂度为O(N^(K-1))。
class Solution: def fourSum(self, nums: List[int], target: int) -> List[List[int]]: if not nums or len(nums) < 4: return [] result = [] nums.sort() cur = 0 while cur < len(nums) - 3: n = nums[cur] for x in self._three_sum(nums[cur + 1:], target - n): result.append([n] + x) cur += 1 # 跳过重复元素 while cur < len(nums) - 3 and nums[cur - 1] == nums[cur]: cur += 1 return result @classmethod def _three_sum(cls, nums: List[int], target: int): """ 类似于three sum的逻辑 :param nums: sorted list :param target: :return: """ i = 0 while i < len(nums) - 2: # 左指针 left = i + 1 # 右指针 right = len(nums) - 1 while left < right: tmp = [nums[i], nums[left], nums[right]] _sum = sum(tmp) if _sum == target: yield tmp right -= 1 left += 1 # 跳过重复元素 while left < right and nums[left] == nums[left - 1]: left += 1 # 跳过重复元素 while left < right and nums[right] == nums[right + 1]: right -= 1 elif _sum > target: right -= 1 else: left += 1 i += 1 # 跳过重复元素 while i < len(nums) - 2 and nums[i - 1] == nums[i]: i += 1
``