色んな解法が考えられそうな問題で、ドツボにハマりやすくて危険な問題だったと思う。落ち着いて整理してシンプルに考える力が問われる。実は、二分探索などは必要ない問題である。

問題概要

$N$ 人の選手 $0, 1, \dots, N-1$ がいる。選手 $i$ の出身国は $C_{i}$ であり、座標 $X_{i}$ の位置に座っている。

各選手 $i$ に対して、

その選手 $i$ とは出身国が異なる選手のうち、
その選手 $i$ との距離が最も近い選手について、

その距離の最小値を求めよ。

制約

$2 \le N \le 3 \times 10^{5}$
$1 \le C_{i}, X_{i} \le 10^{9}$

考えたこと

この問題と似ている問題で、よく知られているのは次のような問題だと思われる。

drken1215.hatenablog.com

つまり、一直線上に $N$ 点が与えられて、 $Q$ 回のクエリが投げられて、各クエリについて座標値が与えられるので、その位置から最も近い点を求めるというものである。この問題は二分探索 (lower_bound()) をすれば解ける典型的な問題であった。

今回の問題も、この過去の問題と同様に、二分探索による解法を考えたくなる。しかし、大きく異なる点が 2 つある。

選手に「色」のような属性があって、同じ色の選手は除外しなければならない
この問題は、クエリ処理問題ではない。つまり、選手のいる座標値以外の座標については考える必要がない

特に後者は見落としがちだと思う。後者に着目して、シンプルに思考する力が問われている。

選手を座標順に並べる

今回、一旦選手を座標順に並べてみよう。具体例として、選手を (座標, 出身国) と表したとき、

(10, 3), (11, 1), (15, 1), (21, 1), (29, 2), (31, 2), (41, 3)

というデータを考えてみる。このデータに対して出身国だけを抽出すると、

3, 1, 1, 1, 2, 2, 3

のようになる。このとき、もし、各要素に対して、

left[i] ← $i$ 番目よりも左の要素であって、出身国が異なるような要素の添字の最大値 (存在しない場合は -1)
right[i] ← $i$ 番目よりも右の要素であって、出身国が異なるような要素の添字の最小値 (存在しない場合は N)

を求めることができれば、この問題は解ける。各要素 $i$ に対して「i 番目の要素と left[i] 番目の要素の座標値の差」と「i 番目の要素と right[i] 番目の要素の座標値の差」の最小値を答えればよい。

そして、left と right の値も、for 文を用いて容易に求められる。たとえば left は次のように求められる。

vector<int> left(N, -1), right(N, N);
for (int i = 1; i < N; ++i) {
    if (v[i].col == v[i-1].col) left[i] = left[i-1];
    else left[i] = i-1;
}

計算量

最初にすべての選手を座標順にソートする部分が最も計算量を要して、 $O(N \log N)$ の計算量となる。

その後の left, right を求めたり、それを用いて答えを求める部分は $O(N)$ の計算量で実行できる。

コード

#include <bits/stdc++.h>
using namespace std;
void chmin(int &a, int b) { if (a > b) a = b; }

// 選手の情報 (もとの添字, 国を表す数値, 座標)
struct Student {
    int id;
    long long col;
    long long pos;
    Student(int i = 0, long long c = 0, long long p = 0) : id(i), col(c), pos(p) {}
};

const int INF = 1LL<<30;
int main() {
    int N;
    cin >> N;
    
    // 選手を座標順にソートする (選手の元の添字も記録しておく)
    vector<Student> v(N);
    for (int i = 0; i < N; ++i) {
        int c, x;
        cin >> c >> x;
        v[i] = Student(i, c, x);
    }
    sort(v.begin(), v.end(), [&](Student a, Student b){return a.pos < b.pos;});

    // 座標順に、各生徒から見て、前後に見たときの最初の「色が違う選手の位置」を求めていく
    vector<int> left(N, -1), right(N, N);
    for (int i = 1; i < N; ++i) {
        if (v[i].col == v[i-1].col) left[i] = left[i-1];
        else left[i] = i-1;
    }
    for (int i = N-2; i >= 0; --i) {
        if (v[i].col == v[i+1].col) right[i] = right[i+1];
        else right[i] = i+1;
    }
        
    // 答えを集計する
    vector<int> res(N, INF);
    for (int i = 0; i < N; ++i) {
        if (right[i] < N) chmin(res[v[i].id], v[right[i]].pos - v[i].pos);
        if (left[i] >= 0) chmin(res[v[i].id], v[i].pos - v[left[i]].pos);
    }
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}