合法三元数量计算

问题描述

小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

问题理解

我们需要找出数组 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。

数据结构选择

为了高效地解决这个问题，我们可以考虑使用哈希表（或字典）来记录每个元素的出现次数。这样可以帮助我们快速计算符合条件的三元组数量。

算法步骤

统计元素频率：首先遍历数组，统计每个元素的出现次数。
计算三元组数量：
- 对于每个元素 x，检查是否存在 x+1 和 x-1 的元素。
- 如果存在，计算以 x 为中心的三元组数量。

具体步骤

统计频率：使用一个哈希表 freq 记录每个元素的出现次数。
计算三元组：
- 遍历哈希表中的每个元素 x，检查 x+1 和 x-1 是否存在。
- 如果存在，计算以 x 为中心的三元组数量。

代码实现

#include <iostream>
#include <vector>
#include <unordered_map>using namespace std;long long solution(vector<int> a) {// 统计元素频率unordered_map<int, int> freq;for (int num : a) {freq[num]++;}long long count = 0;// 计算三元组数量for (auto& [num, f] : freq) {// 检查 num+1 是否存在if (freq.count(num + 1)) {// 计算以 num 为中心的三元组数量// 这里需要计算组合数，具体公式为 C(f, 2) * freq[num + 1]// 其中 C(f, 2) 表示从 f 个元素中选 2 个的组合数count += (long long)f * (f - 1) / 2 * freq[num + 1];}// 检查 num-1 是否存在if (freq.count(num - 1)) {// 计算以 num 为中心的三元组数量// 这里需要计算组合数，具体公式为 C(f, 2) * freq[num - 1]count += (long long)f * (f - 1) / 2 * freq[num - 1];}}return count;
}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;
}