923. 3Sum 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.
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
30
31
32
33
34
35
36
37
| class Solution {
public:
int threeSumMulti(vector<int>& arr, int target) {
long long res = 0, mod = 1e9 + 7;
map<int, long long> m;
for(auto num : arr) {
m[num]++;
}
sort(arr.begin(), arr.end());
arr.erase(unique(arr.begin(),arr.end()),arr.end());
for(int i = 0; i < arr.size() && arr[i] * 3 <= target; i++) {
for(int j = i; j < arr.size() && arr[i] + arr[j] * 2 <= target; j++) {
if(m.count(target - arr[i] - arr[j])) {
if(target == arr[i] * 3) {
if(m[arr[i]] >= 3) {
res = res + m[arr[i]] * (m[arr[i]] - 1) * (m[arr[i]] - 2) / 6;
}
} else if(target == arr[i] + arr[j] * 2) {
if(m[arr[j]] >= 2) {
res = res + m[arr[i]] * m[arr[j]] * (m[arr[j]] - 1) / 2;
}
} else if(i == j){
if(m[arr[i]] >= 2) {
res = res + m[arr[i]] * (m[arr[i]] - 1) / 2 * m[target - arr[i] * 2];
}
} else {
res = res + m[arr[i]] * m[arr[j]] * m[target - arr[i] - arr[j]];
}
res %= mod;
}
}
}
return res;
}
};
|
