3Sum With Multiplicity

Medium

Given an integer arrayA, and an integertarget, return the number of tuples i, j, k such thati < j < kand 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:

• 3 <= A.length <= 3000

• 0 <= A[i] <= 100

• 0 <= target <= 300

Solution

Approach 1: Three Pointer

NSum问题很自然地想到用 N(Two) Pointers

• Sort the array.

• For each i, set T = target - A[i], the remaining target.

• We can try using a two-pointer technique to find A[j] + A[k] == T.

Whenever A[j] + A[k] == T, we should count the multiplicity of A[j] and A[k].

因为有重复元素，我们需要特别小心。在A[j] + A[k] == T时，我们要分两种情况计算multiplicity：

• A[j] == A[k]，这个是说明A[j], A[k]本身都在同一个重复元素序列中，并且一个是头一个是尾，这样根据组合数的计算C(n, 2)，即n个元素取2个，有n (n - 1) / 2种，因此结果 count += (j - k + 1) * (j - k) / 2;

• A[j] != A[k]，这个时候A[j], A[k]不在同一个重复元素序列中，但是本身可能会有重复，因此另设计数变量left, right，分别记录A[j], A[k]有多少次重复，最后计算两者相乘即可，即 count += left * right;

LeetCode Solution

class Solution {
    public int threeSumMulti(int[] A, int target) {
        int MOD = 1_000_000_007;
        long ans = 0;
        Arrays.sort(A);

        for (int i = 0; i < A.length; ++i) {
            // We'll try to find the number of i < j < k
            // with A[j] + A[k] == T, where T = target - A[i].

            // The below is a "two sum with multiplicity".
            int T = target - A[i];
            int j = i+1, k = A.length - 1;

            while (j < k) {
                // These steps proceed as in a typical two-sum.
                if (A[j] + A[k] < T)
                    j++;
                else if (A[j] + A[k] > T)
                    k--;
                else if (A[j] != A[k]) {  // We have A[j] + A[k] == T.
                    // Let's count "left": the number of A[j] == A[j+1] == A[j+2] == ...
                    // And similarly for "right".
                    int left = 1, right = 1;
                    while (j+1 < k && A[j] == A[j+1]) {
                        left++;
                        j++;
                    }
                    while (k-1 > j && A[k] == A[k-1]) {
                        right++;
                        k--;
                    }

                    ans += left * right;
                    ans %= MOD;
                    j++;
                    k--;
                } else {
                    // M = k - j + 1
                    // We contributed M * (M-1) / 2 pairs.
                    ans += (k-j+1) * (k-j) / 2;
                    ans %= MOD;
                    break;
                }
            }
        }

        return (int) ans;
    }
}

My Version (63 ms, 53.40%)

class Solution {
    public int threeSumMulti(int[] A, int target) {
        if (A == null || A.length < 3) {
            return 0;
        }

        Arrays.sort(A);
        int MOD = 1000000007;
        int count = 0;
        int i, j, k;

        for (i = 0; i < A.length - 2; i++) {
            j = i + 1;
            k = A.length - 1;

            int t = target - A[i];

            while (j < k) {
                if (A[j] + A[k] > t) {
                    k--;
                } else if (A[j] + A[k] < t) {
                    j++;
                } else if (A[j] != A[k]) {
                    int left = 1;
                    int right = 1;
                    while (j < k - 1 && A[j] == A[j + 1]) {
                        j++;
                        left++;
                    }
                    while (j + 1 < k && A[k] == A[k - 1]) {
                        k--;
                        right++;
                    }
                    count += left * right;
                    count %= MOD;

                    if (j + 1 == k) {
                        break;
                    } else {
                        j++;
                        k--;
                    }
                } else {
                    count += (k - j) * (k - j + 1) / 2;
                    count %= MOD;
                    break;
                }
            }
        }
        return count;
    }
}

// sorted
// [1,1,1,1,1.5,1.5,2,2,3], target = 4
//  ^       ^    ^       
// targetRest = 4 - 1 = 3

// O(nlogn) + O(n^2) ~ O(n^2)

Time Complexity: O(N^2), where N is the length of A.

Space Complexity: O(1).

Approach 2: Adapt from Three Sum

HashMap - Store number of occurrence (17ms, 73.79%)

class Solution {
    int M = 1000000007;

    public int threeSumMulti(int[] A, int target) {
        Arrays.sort(A);    
        long count = 0;
        Map<Integer, Long> map = new HashMap<>();    

        for (int i = 0; i < A.length; i++) {
            map.put(A[i], 1l + map.getOrDefault(A[i], 0l));
        }

        for (int i = 0; i < A.length; i++) {
            if (i > 0 && A[i] == A[i - 1])
                continue;

            int j = i + 1;
            int k = A.length - 1;                        

            while (j < k) {
                if (A[i] + A[j] + A[k] == target) {
                    if (A[i] != A[j] && A[j] != A[k]) {
                        count += map.get(A[i]) * map.get(A[j]) * map.get(A[k]); 
                    }
                    else if (A[j] != A[k])
                        count += map.get(A[k]) * map.get(A[j]) * (map.get(A[j]) - 1) / 2 % M;
                    else if (A[i] != A[j])
                        count += map.get(A[i]) * map.get(A[j]) * (map.get(A[j]) - 1) / 2 % M;
                    else {
                        count += map.get(A[j]) * (map.get(A[j]) - 1) * (map.get(A[j]) - 2) / 6 % M;      
                        return (int) count;
                    }
                    j++;
                    k--;
                    while (j < k && A[j] == A[j - 1])
                        j++;                    

                    while (j < k && A[k] == A[k + 1])
                        k--;                

                } else if (A[i] + A[j] + A[k] < target) {
                    j++;
                    while (j < k && A[j] == A[j - 1])
                        j++;                    
                } else if (A[i] + A[j] + A[k] > target) {                   
                    k--;
                    while (j < k && A[k] == A[k + 1])
                        k--;                
                }                   
            }
        }
        return (int)(count % M);
    }
}

Approach 3: Think outside of the box - build a map for counting different sums of two numbers

外层循环i，内层循环j，内层循环不断更新A[i] + A[j] two sum对应的计数，外层i在下一个循环时，则相当于第三个index k，若找到target - A[i]就更新结果。

Time - O(N^2), Space - O (N^2)（590 ms, faster than 9.71%）

class Solution {
    public int threeSumMulti(int[] A, int target) {
        Map<Integer, Integer> map = new HashMap<>();

        int res = 0;
        int mod = 1000000007;
        for (int i = 0; i < A.length; i++) {
            res = (res + map.getOrDefault(target - A[i], 0)) % mod;

            for (int j = 0; j < i; j++) {
                int temp = A[i] + A[j];
                map.put(temp, map.getOrDefault(temp, 0) + 1);
            }
        }
        return res;
    }
}

Reference

https://leetcode.com/problems/3sum-with-multiplicity/solution/