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

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

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

これかなりいい問題だと思う!!!
テスターをしていて、最初は 3 人じゃなくて  K 人だったけど、3 人にしたことでいい感じに 300 点問題になった!

問題概要

 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;
}