問題概要

$N$ 個の整数 $A_{1}, A_{2}, \dots, A_{K}$ が与えられる。これらの整数から異なる 2 個を選ぶ方法のうち、2 個の値の差が $K$ 以下であるものの個数を求めよ。

制約

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

解法 (1)：しゃくとり法

鉄則本の問題なので、鉄則本を参照

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

int main() {
    long long N, K;
    cin >> N >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];
    
    // しゃくとり法
    long long res = 0;
    vector<int> right(N, 0);
    for (int l = 0; l < N; ++l) {
        // 前回地点からスタートする
        if (l > 0) right[l] = right[l - 1];
        
        // right を右へと動かす
        while (right[l] < N && A[right[l]] - A[l] <= K) {
            ++right[l];
        }
        --right[l];
        
        res += right[l] - l;
    }
    cout << res << endl;
}

解法 (2)：変数固定 + 二分探索

もう 1 つ、 $A_{i}$ を固定して、 $A_{j} \le A_{i} + K$ を満たす最大の $j$ を二分探索 (lower_bound() など) によって求める方法を考える。

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

int main() {
    long long N, K;
    cin >> N >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];
    
    long long res = 0;
    for (int i = 0; i < N; ++i) {
        long long j = upper_bound(A.begin(), A.end(), A[i] + K) - A.begin();
        --j;
        res += j - i;
    }
    cout << res << endl;
}