Leetcode 923 3Sum With Multiplicity

发表于 |

给定一个整数数组 A，以及一个整数 target 作为目标值，返回满足 i < j < k 且 A[i] + A[j] + A[k] == target 的元组 i, j, k 的数量。
由于结果会非常大，请返回 结果除以 10^9 + 7 的余数。

示例 1：
输入：A = [1,1,2,2,3,3,4,4,5,5], target = 8
输出：20
解释：
按值枚举（A[i]，A[j]，A[k]）：
(1, 2, 5) 出现 8 次；
(1, 3, 4) 出现 8 次；
(2, 2, 4) 出现 2 次；
(2, 3, 3) 出现 2 次。

示例 2：
输入：A = [1,1,2,2,2,2], target = 5
输出：12
解释：
A[i] = 1，A[j] = A[k] = 2 出现 12 次：
我们从 [1,1] 中选择一个 1，有 2 种情况，
从 [2,2,2,2] 中选出两个 2，有 6 种情况。

提示：

  1. 3 <= A.length <= 3000
  2. 0 <= A[i] <= 100
  3. 0 <= target <= 300

题意分析：
在整数数组A中找三个数使三个数的和等于给定值，A中包含重复元素。
A的长度不超过3000，理论上n^2的复杂度是可以被接受的。

思路分析：
我们很轻易的就能想到固定住一个数a[i]，然后在剩下的元素找和为target - a[i]的结果，就把问题转换成了一个twoSum问题。

思路没有问题，但是在这里如何去写twoSum()呢
通常我们有两种方法，一种是使用空间vis去保存已经遍历过的数，然后观察target-A[i]是否存在于vis中.

1
2
3
4
5
6
7
8
9
10
11
# O(n) time, O(n) space
import collections
def twoSum(A,target):
    if target < 0: return 0
    d = collections.Counter()
    res = 0
    for val in A:
        if target - val in d:
            res += d[target-val]
        d[val] += 1
    return res

另一种是先排序，然后使用双指针法遍历，但是因为存在重复数字，所以写起来并不是那么简单。
当然实际排序在3Sum函数中已经排好序了。这里函数是对已经排好序的数组使用双指针法的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# O(n) time, O(1) space, 2608ms
def twoSum(A,target):
    if target < 0: return 0
    i,j = 0,len(A)-1
    ans = 0
    while i < j:
        while i < j and A[i] + A[j] < target: i += 1
        while i < j and A[i] + A[j] > target: j -= 1
        # 当相等之后开始分析重复数字
        if A[i] + A[j] == target:
            if A[i] == A[j]:
                ans += (j-i+1)*(j-i) // 2
                break
            else:
                l,r = 1,1
                while i+1 < j and A[i+1] == A[i]: l += 1;i += 1
                while i < j-1 and A[j-1] == A[j]: r += 1;j -= 1
                ans += l * r
                i += 1
                j -= 1
    return ans

def threeSumMulti(A, target):
    res = 0
    mod = 10**9+7
    A.sort()
    for i in range(len(A)):
        res += twoSum(A[i+1:], target-A[i])
    return res % mod

第一种方法简单，但是却会TLE，首先Counter()函数的构造是很花时间的；就算把Counter()函数换成了d = [0] * 100(因为A[i]的取值不超过100)，还是TLE。

因为我自己代码逻辑的问题，每次twoSum都用重新开空间，确实会有很大的重复计算，但是我想每次重开的空间d，其实数量并不大，也就不到100，A的长度不超过3000，总空间也不会超过300000啊，感觉应该是能AC的。

然后我写了个O(1) space的解法，最后勉强能通过，我猜测可能是python本身语言特性的原因导致tle。

当然为了验证这个猜测，我也仿照discuss里面别人c++的思路来写python代码，但只要时间复杂度为O(n^2)并涉及到空间的都会TLE。

唯独这一种解法，在使用单纯的数组而非Counter()函数后，是勉强可以通过的，运行时间和我的解法差不多。(所以我脑子还是不够灵性，像这样计算前面的twoSum就不用每次都重开数组了)

1
2
3
4
5
6
7
8
9
10
11
# O(n^2) time, O(300) space 2744ms
class Solution(object):
    def threeSumMulti(self, A, target):
        d = [0] * 300
        res = 0
        for i in range(len(A)):
            res += d[target-A[i]] if target-A[i] >=0 else 0
            res %= (10**9+7)
            for j in range(i):
                d[A[i] + A[j]] += 1
        return res % (10**9+7)