実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 180

問題概要

数直線上に $N$ 匹の蟻がおり、蟻 $i$ は始め相異なる座標 $X_i$ にいて正負どちらかの方向を向いている。蟻が向いている向きは0 (負方向) , 1 (正方向) からなる長さ $N$ の文字列 $S$ で与えられる。現在を時刻 $0$ とし、時刻 $T + 0.1$ まで全ての蟻が各々の向いている方向に単位時間あたり $1$ の速さで動く。時間 $T + 0.1$ の間に蟻 $i$ と $j$ がすれ違う整数組 $(i, j) \: (1 \leq i \lt j \leq N)$ の個数を求めよ。

制約

$N, T, X_i$ は整数。
$2 \leq N \leq 2 \times 10^5$
$1 \leq T \leq 10^9$
$-10^9 \leq X_i \leq 10^9$

考察

正方向を向いている蟻に注目し、時間 $T + 0.1$ の間に負方向を向いた蟻と何匹すれ違うかを考えよう。

ここで、単位時間あたりに正方向を向いた蟻と負方向を向いた蟻の距離は $2$ 縮まることから、負方向を向いた蟻を数直線上に固定して、正方向を向いた蟻を単位時間あたり $2$ の速さで移動させても状況は変わらない。

したがって、正方向を向いた各蟻について、その蟻の位置から正方向に $2(T+0.1)$ までの範囲内にいる負方向を向いた蟻の数を数えて、それを足し合わせれば答えが求められることになる。

さて、この蟻の数を数えるにあたって、愚直にループを用いると各蟻に対して $O(N)$ かかってしまい実行時間制限に間に合わない。そこで、負方向を向いた蟻の位置をソートしておくことで、二分探索により対数時間まで計算量を抑えることができる。

コード

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

// ======================================== //

int main()
{
    int N, T;
    cin >> N >> T;
    string S;
    cin >> S;
    vector<ll> left, right;
    rep(i, 0, N)
    {
        ll x;
        cin >> x;
        if (S[i] == '0')
        {
            left.push_back(x);
        }
        else
        {
            right.push_back(x);
        }
    }

    sort(all(left));

    ll ans = 0;
    for (auto &&x : right)
    {
        ans += upper_bound(all(left), x + 2 * (T + 0.1)) - lower_bound(all(left), x);
    }

    cout << ans << endl;
}