【codeFlyer】C - 徒歩圏内

問題概要

N個の都市がそれぞれ座標X_iにある. 2つの都市間の距離がD以下であれば徒歩で, そうでなければ電車で移動する.

このとき, 3つの都市の組(i, j, k)であり, 以下の条件を満たすものの個数を求めよ.

• i < j < k
• 都市iと都市j, 都市jと都市kを徒歩で移動する.
• 都市iと都市kを電車で移動する.

制約

• 3 <= N <= 10⁵
• 0 <= X_i <= 10⁹
• 0 <= D <= 10⁹

考察

二分探索がほのかに香ってきたのでその方針で考える.

都市iから都市j, 都市jから都市kまで徒歩で移動する場合の数から, 都市iから都市kまで徒歩で移動する場合の数を引けば良さそう. 前者をA, 後者をBと置いて求め方を考えていく.

Aについて

都市iを固定する. 都市iから徒歩で行ける都市j全てについて, さらにその都市jから徒歩で行ける都市kの数を数えれば良い.

それぞれの都市から徒歩で行ける都市の数をあらかじめ数えておく. これはそれぞれの都市について, どの都市まで徒歩で行けるかを二分探索することで求められる(O(NlogN)). (変数名はwalk)

次に, それぞれの都市から徒歩で行ける都市について, walkの和を順に足していきたい. 徒歩で行ける都市はもちろん並んでいるので, walkの累積和を取ることでオーダーを落とせる.

よって, どの都市まで徒歩で行けるかを二分探索で調べ, そのwalkの区間和を累積和で求めることによりAが求められた(O(NlogN)).

Bについて

都市iを固定する. そこから徒歩で行ける都市の数を二分探索で求める. それらの都市から適当に2つ選べばBの条件を必ず満たすので, Bが求められた(O(NlogN)).

ACコード(C++🔆)

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define REP(i,n) for(int (i)=0; (i)<(n); ++(i))
#define LEQ(v,x) (upper_bound((v).begin(), (v).end(), x) - (v).begin())
 
int N, D;
ll res;
vector<ll> X, walk, walk_acc;
 
ll comb(ll num) {
    return max<ll>(0, num * (num - 1) / 2);
}

int main() {

    cin >> N >> D;
    X.resize(N);
    REP(i, N) cin >> X[i];

    REP(i, N) walk.push_back(LEQ(X, X[i] + D) - i - 1);

    walk_acc.push_back(0);
    REP(i, N) walk_acc.push_back(walk_acc[i] + walk[i]);
 
    REP(i, N) {
        res += walk_acc[min<ll>(N, LEQ(X, X[i] + D))] - walk_acc[i + 1];
        res -= comb(LEQ(X, X[i] + D) - i - 1);
    }

    cout << res << endl;

}