0923. 3Sum With Multiplicity (M)

题目

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 10^9 + 7.

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.

Constraints:

• 3 <= arr.length <= 3000
• 0 <= arr[i] <= 100
• 0 <= target <= 300

题意

在指定数组中取三个数，使它们的和正好为目标值，求这样的三元组合的个数。

思路

类似3sum，每次固定i，使用双指针j和k遍历i之后的数组。关键在于 arr[i]+arr[j]+arr[k] == target 时要进行分类讨论：1. 如果 arr[j] == arr[k]，说明arr[j~k]全一样，那么需要统计从中任选2个位置能得到的二元对的个数；2. 如果 arr[j] != arr[k]，那么需要找到从j开始向后等于arr[j]的元素的长度，从k开始向前等于arr[k]的元素的长度，相乘得到对应的二元对的个数。

代码实现

Java

class Solution {
    public int threeSumMulti(int[] arr, int target) {
        int ans = 0;
        Arrays.sort(arr);
        for (int i = 0; i < arr.length - 2; i++) {
            int j = i + 1, k = arr.length - 1;
            while (j < k) {
                if (arr[i] + arr[j] + arr[k] > target) {
                    k--;
                } else if (arr[i] + arr[j] + arr[k] < target) {
                    j++;
                } else if (arr[j] != arr[k]){
                    int left = 1, right = 1;
                    while (j + 1 < k && arr[j + 1] == arr[j]) {
                        j++;
                        left++;
                    }
                    while (k - 1 > j && arr[k - 1] == arr[k]) {
                        k--;
                        right++;
                    }
                    ans = (ans + left * right) % 1000000007;
                    j++;
                    k--;
                } else {
                    ans = (ans + (k - j + 1) * (k - j) / 2) % 1000000007;
                    break;
                }
            }
        }
        return ans;
    }
}