codeFlyer （bitFlyer Programming Contest） C - 徒歩圏内 (400)

C - 徒歩圏内

概要

長さNの非負整数列X(i < j ⇒ X[i] < X[j])と非負整数Dが与えられる。下記の条件を満たす三つ組の整数(i, j, k) (i < j < k)の個数を求めよ。

• X[j] - X[i] ≦ D かつ X[k] - X[j] ≦ D かつ X[k] - X[i] > D

制約

3 ≦ N ≦ 10^5
0 ≦ X[i], D ≦ 10^9

方針

このままだと数えづらいので、下記2つの数を数えることにする。
A：X[j] - X[i] ≦ D かつ X[k] - X[j] ≦ D なる(i, j, k) (i < j < k)の個数
B：X[k] - X[i] ≦ D なる(i, j, k) (i < j < k)の個数
このとき、求める数はA - B となる。

Aの個数は下記のように求める。
①iを固定する。
②各iについてX[j] - X[i] ≦ Dなるjを固定する。
③各jについてX[j] - X[i] ≦ Dなるkを数える。
これは、y[i] = #{ j > i | X[j] - X[i] ≦ D} を定め、y[i]の累積和としてz[i]を準備しておくことで計算を効率化できる。

Bの個数は j = max({j' > i | X[j'] - X[i] ≦ D}) とすると、各iについてi < j' < k ≦ j なる(j', k)の個数であるため、(j - i - 1) * (j - i) / 2個である。

解答

Submission #8210265 - codeFlyer （bitFlyer Programming Contest）

#include <bits/stdc++.h>
#define rep(i,n) for(int i=(0);i<(n);i++)
 
using namespace std;
 
typedef long long ll;
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	ll n, d;
	cin >> n >> d;
 
	vector<ll> x(n);
	rep(i, n) cin >> x[i];
 
	vector<ll> y(n);
	rep(i, n){
		auto itr = lower_bound(x.begin(), x.end(), x[i] + d + 1);
		y[i] = distance(x.begin(), itr) - i - 1;
	}
 
	// rep(i, n) cout << y[i] << (i == n - 1 ? "\n" : " ");
 
	vector<ll> z(n, 0);
	z[0] = y[0];
 
	rep(i, n - 1) z[i+1] = z[i] + y[i+1];
 
	ll ans1 = 0, ans2 = 0;
	for(ll i = 0; i < n; i++) {
		auto itr = lower_bound(x.begin(), x.end(), x[i] + d + 1);
		if(itr == x.begin()) continue;
		itr--;
 
		ll j = distance(x.begin(), itr);
		// cout << i << "," << j << endl;
		ans1 += z[j] - z[i];
		ans2 += (j - i - 1) * (j - i) / 2;
 
	}
 
	// cout << ans1 << "," << ans2 << "," << ans1 - ans2 << endl;
	cout << ans1 - ans2 << endl;
 
}