[Codeforces] Codeforces Round #690 (Div. 3) E1. Close Tuples (easy version)

2020年12月16日2021年6月19日

問題

方針

\[ \max(a_i, a_j, a_k) – \min(a_i, a_j, a_k) \leq 2 \tag{1} \]

\( (1)\) を満たすような \( a_i, a_j, a_k \) は、\( i, j, k \) を交換することによって、\( i < j < k \) を満たすことができます。つまり、\( (1)\) を満たすような \( a_i, a_j, a_k \) が存在すれば、\( i < j < k \) を満たすような選び方ができるということです。したがって、配列の並びを変えても答えは変わらないので、\( a_i \) を昇順にして考えます。

ここで、自然数 \( l, r \) について \( l \) を固定したとき、\( a_r – a_l \leq 2 \) を満たす最大の \( r \) を求めます。 このとき、\(a_l \) を含み、\( (1) \) と \( i < j < k \) を満たすような組み合わせは、

\[ \dfrac{(r – l – 1)(r – l)}{2}\]

通りあります。したがって、\( l, r \) について全探索します。\( l + 1 \) について調べるとき、\( r \) の値は、\( l \) に対応する \( r \) 以上となるので、しゃくとり法を用いて高速化します。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

int n;
int a[200000];

void solve() {
    if (n <= 2) {
        cout << "0\n";
        return;
    }
    sort(a, a + n);
    ll cnt = 0;
    int r = 2;
    for (int l = 0; l < n - 1; l++) {
        while (r + 1 < n) {
            if (a[r + 1] - a[l] <= 2) r++;
            else break;
        }
        if (r + 1 < n) {
            if (a[r + 1] - a[l] <= 2) r++;
        }
        if (l > r) {
            r = l + 1;
            continue;
        }
        if (a[r] - a[l] <= 2) {
            ll d = r - l;
            cnt += d * (d - 1) / 2;
        }
    }
    cout << cnt << "\n";
}

int main() {
    int t;
    cin >> t;
    while (t--) {
        cin >> n;
        for (int i = 0; i < n; i++) {
            cin >> a[i];
        }
        solve();
    }
    return 0;
}