Three Sum With Multiplicity

Given an integer array arr, and an integer target, return the number of tuples i, j, k such that i < j < k and arr[i] + arr[j] + arr[k] == target.

As the answer can be very large, return it modulo 109 + 7.

Example

Example 1:

Input: arr = [1,1,2,2,3,3,4,4,5,5], target = 8
Output: 20
Explanation: 
Enumerating by the values (arr[i], arr[j], arr[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: arr = [1,1,2,2,2,2], target = 5
Output: 12
Explanation: 
arr[i] = 1, arr[j] = arr[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

固定一个数，做2sum

The tricky parts are two:

• if the two sum all the elements are the same, then we count the length * (length - 1) / 2 like [2,2,2,2] is 6 ways

• otherwise, we track the duplicate left or right elements and get the product of them

Code

class Solution:
    def threeSumMulti(self, nums: List[int], target: int) -> int:
        MOD = 10**9 + 7
        nums.sort()
        
        res = 0
        for i in range(len(nums) - 2):
            left = i + 1
            right = len(nums) - 1
            add = target - nums[i]
            while left < right:
                if nums[left] + nums[right] < add:
                    left += 1
                elif nums[left] + nums[right] > add:
                    right -= 1
                else:
                    l = r = 1
                    while left < right and nums[left] == nums[left + 1]:
                        left += 1
                        l += 1
                    while left < right and nums[right] == nums[right - 1]:
                        right -= 1
                        r += 1
                    if left != right: 
                        res += l * r
                    else:
                        res += l * (l - 1) // 2
                    left += 1
                    right -= 1
        
        return res % MOD