問題概要

$N$ 人がいて、それぞれレーティングは $a_{1}, \dots, a_{N}$ となっている。

$N$ 人の中から 3 人を選ぶ方法のうち、3 人のレーティングの最大値と最小値との差が $D$ 以下であるようなものが何通りあるか、求めよ。

制約

$3 \le N \le 10^{5}$

解法

この手の問題は最初に $N$ 人をレーティングの小さい順にソートしてしまうと考えやすくなる。

そして「3 人のうちレーティングが最小の人」を誰にするのかを固定して考えることにする。 $i$ 人目をレーティング最小の人としたとき、

$a_{j} \le a_{i} + D$ を満たす最大の $j$ をとる
このとき $i+1, i+2, \dots, j$ の中からどのように 2 人選んでもよい
その選び方は ${}_{j-i}{\rm C}_{2} = \frac{1}{2}(j-i)(j-i-1)$ 通りある

という風になる。あとは各 $i$ に対して、 $a_{j} \le a_{i} + D$ を満たす最大の $j$ が求められれば OK。これはしゃくとり法 (か lower_bound) を使えば OK。

全体として $O(N \log N)$ の計算量で解ける。

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

int main() {
    int N, D;
    cin >> N >> D;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    sort(a.begin(), a.end());

    long long res = 0;
    int j = 0;
    for (int i = 0; i < N; ++i) {
        while (j < N && a[j] <= a[i] + D) ++j;
        --j;
        res += (long long)(j-i) * (j-i-1) / 2;
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

CPSCO2019 Session3 C - Make a Team (300 点設定)

問題概要

制約

解法