LeetCode之923. 3Sum With Multiplicity

题目链接：https://leetcode.com/problems/3sum-with-multiplicity/
题目描述：
Given an integer array A, and an integer target, return the number of tuples i, j, k such that i < j < k and A[i] + A[j] + A[k] == target.
As the answer can be very large, return it modulo 10^9 + 7.
Example 1

Input: A = [1,1,2,2,3,3,4,4,5,5], target = 8
Output: 20
Explanation: 
Enumerating by the values (A[i], A[j], A[k]):
(1, 2, 5) occurs 8 times;
(1, 3, 4) occurs 8 times;
(2, 2, 4) occurs 2 times;
(2, 3, 3) occurs 2 times.

Example 2

Input: A = [1,1,2,2,2,2], target = 5
Output: 12
Explanation: 
A[i] = 1, A[j] = A[k] = 2 occurs 12 times:
We choose one 1 from [1,1] in 2 ways,
and two 2s from [2,2,2,2] in 6 ways.

Note

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

这道题和之前3SUM的题目很像，只不过这里允许有重复的答案，而且可能存在大量重复的答案，所以最后要对一个大数求余。按照 3SUM 的解法，首要要对数组进行排序，但是在做这道题的时候，刚开始一直卡在题中描述的such that i < j < k and A[i] + A[j] + A[k] == target，以为不能对原始数组进行排序，因为排了序后，位置就变了。后来看了一些参考答案之后才想清楚，只要找出和为 target 的这三个数就行了，至于 index，必然是满足递增关系的，就好比随机生成三个不一样的 index，如 5,4,3 ，自然满足3<4<5的关系，只不过索引上对应的数字在前在后罢了。对应到这道题里，假如原来有一组答案是 i<j<k, A[i]+A[j]+A[k]=target，那么排序后，这三个数还是一组答案，只不过在排序后的数组中就可能对应着 A'[3]+A'[0]+A'[4]=target了，原来的 A[i] 是现在的 A’[3] ，A[j] 是现在的 A’[0] ，A[k] 是现在的 A’[4] 。

排序的困扰解决之后，那么和 3SUM 不同的地方就是如何解决重复问题了，基本思路还是一样，先固定一个数，然后使用双指针寻找后面两个数，只不过双指针找到满足条件的两个数后，还要注意重复的问题，原来是有重复就跳过了，现在需要统计有几个重复的数，这里就有2种情况了：

1. 刚开始的A[left]==A[right], 因为数组先排过序了，所以中间这一段也都是相等数字，假设从左到右一共有6个数字，那满足条件的组合共有 5+4+3+2+1 个，使用高中的求和公式，那就是 (r-l)*(r-l+1)/2 个；
2. 刚开始的A[left]!=A[right]，假若左边一共有 l 个相同数字，右边共有 r个相同数字，则满足条件的组合共有 l*r 个

具体代码如下：

def threeSumMulti1(A, target):
    # 排序
    A = sorted(A)
    cnt = 0
    MOD = pow(10,9)+7
    for i in range(len(A)-2):
        if A[i] > target: break

        diff = target - A[i]
        j = i + 1
        k = len(A) - 1
        while j < k:
            if A[j] + A[k] < diff:
                j += 1
            elif A[j] + A[k] > diff:
                k -= 1
            else:
                # sum to target
                # left记录重复数字个数，right记录右边重复数字个数
                left = right = 1
                while(j + left < k and A[j] == A[j + left]): left += 1
                # j+left=k, or A[j]!=A[j+left]
                # 为什么是 >=，因为上一步 j+left 停止的位置可以是A[j]!=A[j+left]，但这个位置可以是和 A[k]相等的，所以是 >=，比如 2,2,2,3,3 这个例子
                while(k - right >= j + left and A[k] == A[k - right]): right += 1
                # k-right=j+left, or A[k]!=A[k-right]

                if A[j]==A[k]:
                    # 满足条件，并且全相等的数字，假设有6个，则一共新增 5+4+3+2+1 个答案
                    cnt += int((k - j) * (k - j + 1) / 2)
                else:
                    cnt += left * right
                j += left
                k -= right

    return cnt % MOD

上面解法在 LeetCode 中提交是超时的，但是换成 C++ 的代码就可以 AC 了，看来 Python 语言刷题还是在速度上略输一筹。。。

上面的解法是按照 3SUM 的思路改的，这题还可以换个思路，就是不用排序，先遍历下数组，并用字典存储数组中每个数字出现的次数。然后用两个 for 循环遍历字典的 key，找出前两个数 i,j，这样可以求出第三个数 k=target-i-j，如果 k 不在字典中，则直接跳过，如果 k 在字典中，则会有以下三种情况：

1. i==j==k：如 5个2，target=6，这5个2对应5个不同的 index，任意取3个都会满足 index 相互递增，那就变成了组合问题了，从n个里任取3个组合, 共有 d[i]*(d[i]-1)*(d[i]-2)/6个；
2. i==j&j!=k，如5个2，2个3，target=7，对应7个index，这时候就变成了从5个2里任取2个index，然后每一个组合都可以和3的一个index拼起来,共有 d[i]*(d[i]-1)/2*d[k]个；
3. i<j<k: 这种情况最简单，如2个1，3个2，2个3，target=6，那就是从1的index里任取1个，2的index里任取1个，3的index里任取1个，拼在一起都满足题中条件，所以共有2*3*2=12个组合，也就是d[i]*d[j]*d[k]个；

为什么是这三种情况：前两种很好理解，两个 for 循环分别遍历，就会出现 i==j 的情况，此时和 k 的关系只有 j==k or j!=k 两种。当 i!=j时，首先不能直接用 i!=j判断，因为会出现同一个答案统计多次的情况，那么 i 和 j 就要有大小关系，才能避免重复，剩下的和 k 的关系只能是 i<j<k or i>j>k，换成 i>j>k也是可以的，但是其他的关系，也会出现重复统计的问题。具体代码如下：

def threeSumMulti12(A, target):
    num_cnt = {}
    result = 0
    MOD = pow(10,9)+7
    for a in A:
        num_cnt[a] = num_cnt.get(a, 0) + 1

    # 0<=A[i]<=100
    for i in num_cnt.keys():
        for j in num_cnt.keys():
            k = target - i - j
            if k not in num_cnt:continue

            if i == j and j==k:
                result += int(num_cnt[i] * (num_cnt[j]-1) * (num_cnt[k]-2) / 6)
            elif i == j and j!=k :
                result += int(num_cnt[i] *(num_cnt[i]-1)/2 * num_cnt[k])
            elif i < j < k:    
                result += num_cnt[i]*num_cnt[j]*num_cnt[k]
            else:
                pass

    return result % MOD

换成下面的写法可能更好理解，因为每个数字的范围是 [0,100]，所以也可以写成下面这样：

def threeSumMulti13(A, target):
    num_cnt = {}
    result = 0
    MOD = pow(10,9)+7
    for a in A:
        num_cnt[a] = num_cnt.get(a, 0) + 1

    # 0<=A[i]<=100
    for i in range(101):
        if i not in num_cnt: continue
        for j in range(i, 101):
            if j not in num_cnt:continue
            k = target - i - j
            if k not in num_cnt:continue

            if i == j and j==k:
                result += int(num_cnt[i] * (num_cnt[j]-1) * (num_cnt[k]-2) / 6)
            elif i == j and j!=k :
                result += int(num_cnt[i] * (num_cnt[i]-1) / 2 * num_cnt[k])
            elif i < j < k:    
                result += num_cnt[i] * num_cnt[j] * num_cnt[k]
            else:
                pass

    return result % MOD

这个写法就能明显的看出来第三种情况是 i<j<k了，时间负责度就是 O(N+101*101)，其中 101 是因为数组中数字的范围是 0-100。

参考资料：

• https://www.cnblogs.com/grandyang/p/11818151.html
• https://leetcode.com/problems/3sum-with-multiplicity/discuss/181131/C%2B%2BJavaPython-O(N-%2B-101-*-101)

本文地址：https://blog.csdn.net/tiangcs/article/details/107300264