Python哈希表一招解决nSum问题
自古各门各家武学都存在套路,正所谓以不变应万变,就在于临战之时,可以一招制敌。有的招数可能出奇制胜,但是最稳定的方式一定是多次训练的套路,它不一定能让你解决所有的问题,但是它足以让你轻松应对一类问题。
$nSum$ 的问题,主要存在大量重复的数使得如果在数组中遍历每个数,再比较查询结果,时间复杂度会超过题目的要求。我们可以采用哈希表的方式,加速查询的过程,同时对遍历的过程,对相同的数或者不满足条件的数适当的跳过,可以有效的提升效率,通过测试用例。
那么先从最简单的两数之和讲起。
1、 Leetcode 1 两数之和
最暴力的方式是从头到尾枚举 $nums$ 中的每一个数,然后再看是否在数组中存在 $j$ 使得 $nums[j] == target - nums[i]$ 且 $j \not= i$,这样的方式,遍历是 $O(n)$ 复杂度,每次循环在数组中查询也是 $O(n)$ 复杂度,总的时间复杂度达到了 $O(n^2)$。代码和运行时间如下:
1 class Solution: 2 def twoSum(self, nums: List[int], target: int) -> List[int]: 3 result = [] 4 for i in range(len(nums)): 5 if nums[i] in nums and target-nums[i] in nums: 6 j = nums.index(target-nums[i]) 7 if i!=j: 8 result.append(i) 9 result.append(j) 10 break 11 return result
这样做显然太过耗时,循环如果不做改变的话(做改变的话可以用双指针法,这里不做过多介绍),那么考虑在查找过程中进行加速。我们注意到哈希表中查找元素的时间是 $O(1)$,因此可以把在数组中查找改为在哈希表中查找。对于这题而言,只要找到了答案就可以返回,不需要找出所有的解,那么可以边遍历边向哈希表中添加元素。添加前,查询是否有满足条件的解,如果满足条件,$return$ 结果就可以。
1 class Solution: 2 def twoSum(self, nums: List[int], target: int) -> List[int]: 3 dic = {} 4 for i,num in enumerate(nums): 5 tmp = target - num # a + b = target, a = num, b = target - num 6 if tmp in dic: # 哈希表中查询是否有解 7 return [i, dic[tmp]] 8 dic[num] = i # 没有解的话就存下当前的数和位置
运行结果如下:
最差的情况下,在线性时间内就可以解决问题
接下来考虑复杂一点的问题
2、 Leetcode 15 三数之和
首先题目要求解集里面不包含重复的元素,那么按照一定的规律找答案,就可以得到不重复的解。可以想到的方法是先进行排序,这样就可以有规律的寻找了。排序以后,每个数也可能有多个重复的,假如每个解里不能包含相同的数字,那么简单的在循环里加上
1 if i > index and nums[i] == nums[i-1]: 2 continue
其中 $index$ 是循环开始的值, 并且下一层循环从 $i + 1$ 开始,就可以保证无重复了。这样的去重方式可以参考我的另一篇文章讲到的第三类问题, https://www.cnblogs.com/HMJIang/p/13575005.html
然而这道题则是每个解里可以包含相同的数字,比如 $[-1, -1, 2]$ 和 $[0, 0, 0]$ 都可以得到和为 $0$,这时候去重就可以用到 $Python$ 中计数哈希表, $Counter$。先统计每个数字出现的次数,再对键值进行排序,每层循环里判断剩余的数字是否够当前的变量选择。
本题解法参考 https://leetcode-cn.com/problems/3sum/solution/ji-shu-zi-dian-jian-zhi-you-hua-fei-pai-xu-shuang-/
代码如下:
1 from collections import Counter 2 class Solution: 3 def threeSum(self, nums: List[int]) -> List[List[int]]: 4 res = [] 5 dic = Counter(nums) # Counter可以统计数组每个元素的个数 6 hash_nums = sorted(dic.keys()) #对键值进行排序 7 for i, a in enumerate(hash_nums): 8 dic[a] -= 1 # a已经取走了一个数字,字典里对应位置 -1 9 for b in hash_nums[i:]: 10 if dic[b] < 1: # b从i开始遍历,i也是当前a的位置,如果减去1以后b不够选了,跳过这一个位置 11 continue 12 c = -(a + b) 13 if c < b: #有序的查找,如果c都比b小,之后b再增大,肯定c更小,那么就跳出,防止重复 14 break 15 # 再判断c和b的关系,如果相等,那就需要dic[c]至少为2,才够选,如果不等,只要有,就够选了 16 if (c > b and dic[c] > 0) or (c == b and dic[c] > 1): 17 res.append([a, b, c]) 18 return res
时间复杂度 $O(n^2)$
空间复杂度 $O(n)$
提交结果:
有了这样的经验以后,我们可以用已有的套路看更复杂的四数之和
3、Leetcode 18 四数之和
最外层循环遍历到什么位置,就在对应位置上 $-1$,接下来内层循环里也把选择的数 $-1$,方便后面进行判断,只要不够选了,就 $continue$ 跳过这一次循环,如果最终的 $d$ 比 $c$ 还大,依旧 $break$ 掉,和三数之和的差别在于,第二个数选择的时候要 $-1$,最内层循环结束以后还要 $+1$,因为之后最外层的 $a$ 也会再遍历到这个位置。
代码如下:
1 from collections import Counter 2 class Solution: 3 def fourSum(self, nums: List[int], target: int) -> List[List[int]]: 4 res = [] 5 dic = Counter(nums) #对每个数出现的次数进行统计 6 arr = sorted(dic.keys()) #排序键值 7 for i, a in enumerate(arr): 8 dic[a] -= 1 #a用掉了一次,而且a的位置之后不会再遍历到了,不需要加回 9 for j, b in enumerate(arr[i:]): #从arr[i]开始找b的值 10 if dic[b] < 1: #b可能等于a,判断一下,如果dic[b]不够1个,跳过这次循环 11 continue 12 dic[b] -= 1 13 for c in arr[i+j:]: #从arr[i+j]开始找c的值,注意上一层循环枚举j以后,需要再加最外层的i 14 if dic[c] < 1: #同上层循环b的判断 15 continue 16 d = target - (a + b + c) 17 if d < c: #因为是非递减顺序,如果d小于c,就直接跳出,这样就可以避免重复 18 break 19 if (d == c and dic[d] > 1) or (d > c and dic[d] > 0): 20 res.append([a, b, c, d]) 21 dic[b] += 1 #b现在所处的位置,之后a还会遍历到,因此需要加回1 22 return res
时间复杂度 $O(n^3)$
空间复杂度 $O(n)$
提交结果:
以此可以类推到更多数字的和。最外层循环每选到一个位置以后,都 $-1$,内层的循环也选到一个位置 $-1$,在更内层的循环结束以后 $+1$ 就可以。最内层转化为两数之间的大小关系的比较和查询哈希表是否有满足条件的值。
最后看一个变种问题
Leetcode 1577 数的平方等于两数乘积的方法数
这一题如果暴力求解,必然超时,那么就需要一些优化策略。两个数组里可能会存在很多相同的数,它们仅仅是位置不同,找到的 $j$, $k$ 的结果却一样,比如 $nums1 = [1,1,1,1]$,$nums2 = [1,1,1,1,1,1]$,$nums1$ 中每个数的平方,都等于 $nums2$ 中任意两个不同位置的数的乘积,我们没有必要对每个相同的 $nums1$ 中的数都找一遍 $nums2$ 中所有的数,这就又回到了 $nSums$ 问题,可以想到的去重的方式是哈希表。
这里的技巧在于如果对于每个平方数去找是否存在两个数和它相等,每个平方数遍历的时间是 $O(n)$, 再找两个数,如果要达到 $O(m)$ 复杂度,就应该考虑双指针的方式,然后需要各种比较,代码相对复杂,容易出错。如果换个思路,从右向左找,对于每个数字的乘积,都在哈希表里找是否存在相应的平方数,那么时间复杂度就是 $两次遍历数组的时间复杂度 × 哈希表查找的时间复杂度$,由于哈希表查找是 $O(1)$,最终等于两次遍历数组的时间复杂度。
代码如下:
1 from collections import Counter 2 class Solution: 3 def numTriplets(self, nums1: List[int], nums2: List[int]) -> int: 4 square1 = Counter([i*i for i in nums1]) 5 square2 = Counter([i*i for i in nums2]) 6 res = 0 7 for i in range(len(nums2)): 8 for j in range(i+1, len(nums2)): 9 tmp = nums2[i] * nums2[j] # 可以用tmp存一下两数之积,避免后面字典查询的时候再次重复计算键值 10 if tmp in square1: 11 res += square1[tmp] 12 for i in range(len(nums1)): 13 for j in range(i+1, len(nums1)): 14 tmp = nums1[i] * nums1[j] # 同理 15 if tmp in square2: 16 res += square2[tmp] 17 return res
时间复杂度 $O(n^2 + m^2)$ 其中 $m$,$n$ 是 $nums1$ 和 $nums2$ 的数组长度
空间复杂度 $O(n+m)$
代码执行结果如下:
之所以说在每次循环中要用 $tmp$ 存一下两数的乘积,因为随着数据量的增加,重复计算两数乘积的代价也是相当大的,如果两次都直接用
1 if nums2[i] * nums2[j] in square1: 2 res += square1[nums2[i] * nums2[j]]
运行时间将会明显提升,提交结果如下: