Codeforces 1462E1 Close Tuples (easy version)

题意

给一个长度为$n$的数组，取$m$个数字，其中最大值最小值相差不大于$k$，问这种方式有多少种。在$easy$ $version$下

$m=3$，$k=2$。

分析

简单版本就该有简单版本的做法，发现如果将$i$视为这$m$个位置中的第一个，只需要考虑最小值是$a_i-2$、$a_i-1$、$a_i$的情况。

最小值是$a_i-2$，由于$a_i$是必取的，所以只需要考虑第三个数是$a_i-2$、$a_i-1$、$a_i$。
最小值是$a_i-1$，由于$a_i$是必取的，所以只需要考虑第三个数是$a_i-1$、$a_i$、$a_i+1$。
最小值是$a_i$，由于$a_i$是必取的，所以只需要考虑第三个数是$a_i$、$a_i+1$、$a_i+2$。

记录下每个数字出现的次数，使用组合数即可。

#pragma GCC optimize(3, "Ofast", "inline")

#include <bits/stdc++.h>

#define start ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define ll long long
#define ull unsigned long long
#define int ll
#define ls st<<1
#define rs st<<1|1
#define pii pair<int,int>
#define rep(z, x, y) for(int z=x;z<=y;++z)
#define repd(z, x, y) for(int z=x;z>=y;--z)
#define com bool operator<(const node &b)const
using namespace std;
mt19937 rnd(chrono::high_resolution_clock::now().time_since_epoch().count());
const int maxn = (ll) 1e6 + 5;
const int mod = 998244353;
const int inf = 0x3f3f3f3f;
int T = 1;
int num[maxn];
int a[maxn];

void solve() {
    int n;
    cin >> n;
    rep(i, 1, n)cin >> a[i], ++num[a[i]];
    int ans = 0;
    rep(i, 1, n) {
        --num[a[i]];
        int now = a[i];
        if (now >= 3) {
            ans += num[now - 2] * (num[now - 2] - 1) / 2 + num[now - 2] * (num[now - 1] + num[now]);
        }
        if (now >= 2) {
            ans += num[now - 1] * (num[now - 1] - 1) / 2 + num[now - 1] * (num[now] + num[now + 1]);
        }
        int tot = 0;
        rep(j, now, now + 2)tot += num[j];
        ans += tot * (tot - 1) / 2;
    }
    cout << ans << '\n';
}

signed main() {
    start;
    cin >> T;
    while (T--)
        solve();
    return 0;
}