问题描述
小C、小U 和小R 三个好朋友喜欢做一些数字谜题。这次他们遇到一个问题，给定一个长度为n的数组a，他们想要找出符合特定条件的三元组 (i, j, k)。具体来说，三元组要满足 0 <= i < j < k < n，并且 max(a[i], a[j], a[k]) - min(a[i], a[j], a[k]) = 1，也就是说，最大值与最小值之差必须为1。
他们决定请你帮忙编写一个程序，计算符合这个条件的三元组数量。
测试样例
样例1：
输入：
a = [2, 2, 3, 1]
输出：2
样例2：
输入：
a = [1, 3, 2, 2, 1]
输出：5
样例3：
输入：
a = [1, 3, 2, 2, 1, 2]
输出：12
#include <iostream>
#include <unordered_map>
#include <vector>
using namespace std;
long long solution(vector<int> a) {
unordered_map<int, long long> freq;
for (int num : a) {
freq[num]++;
}
long long ans = 0;
for (const auto &[x, cnt_x] : freq) {
if (freq.count(x + 1)) {
long long cnt_x1 = freq[x + 1];
// 两个x，一个x+1
ans += (cnt_x * (cnt_x - 1) / 2) * cnt_x1;
// 两个x+1，一个x
ans += (cnt_x1 * (cnt_x1 - 1) / 2) * cnt_x;
}
}
return ans;
}
int main() {
vector<int> a1 = {2, 2, 3, 1};
vector<int> a2 = {1, 3, 2, 2, 1};
vector<int> a3 = {1, 3, 2, 2, 1, 2};
cout << (solution(a1) == 2) << endl;
cout << (solution(a2) == 5) << endl;
cout << (solution(a3) == 12) << endl;
return 0;
}
