値と一緒にインデックスを持たせたいときの一般的なテク

競プロをやっていると、よく ソートされた値にインデックスを持たせたい とか、 RMQの値にインデックスを持たせたい とか、 グラフの辺にインデックスを持たせたい とか、 クエリにインデックスを持たせたい とかいうことがよくあります。

こんなときに使える汎用的なテクニックとは…？

• 値の代わりにインデックスを持つ
• 値を得るときはそのつど値の配列にアクセスする
• ソートとかをするときに、アクセスして比較するような比較関数を与える

ことです。

例を見てみましょう。

ソートされた値にインデックスを持たせる

$N$ 要素の数列

$A_0,A_1,\dots ,A_{N-1}$ があります。

$A_{B_0}\le A_{B_1}\le \cdots \le A_{B_{N-1}}$ となるような

$0,\dots ,N-1$ の順列

$B$ を求めてください。

Before

int N;
cin >> N;
vector<pair<int, int>> A(N);
for(int i = 0; i < N; i++){
    cin >> A[i].first;
    A[i].second = i;
}
sort(A.begin(), A.end());
for(int i = 0; i < N; i++) cout << A[i].second << endl;

After

int N;
cin >> N;
vector<int> A(N), B(N);
for(auto& a : A) cin >> a;
iota(B.begin(), B.end(), 0); // インデックスを持つ　iota が便利
sort(B.begin(), B.end(), [&](int x, int y){ return A[x] < A[y]; }); // Aにアクセスして比較する比較関数
for(auto& b : B) cout << b << endl;

pairを作る必要がない　複数のものをまとめてソートするときかなり有利

priority_queue で値にインデックスを持たせる

Before

int N;
cin >> N;
priority_queue<pair<int, int>> q;
for(int i = 0; i < N; i++){
    int A;
    cin >> A;
    q.emplace(A, i);
}

After

int N;
cin >> N;
vector<int> A(N);
for(auto& a : A) cin >> a;
priority_queue q([&](int x, int y){ return A[x] < A[y]; }, vector<int>());
for(int i = 0; i < N; i++) q.push(i);

C++17 を使うとこうなる
C++17 を使わない場合、

int N;
cin >> N;
vector<int> A(N);
for(auto& a : A) cin >> a;
auto comp = [&](int x, int y){ return A[x] < A[y]; };
priority_queue<int, vector<>, decltype(comp)> q(comp);
for(int i = 0; i < N; i++) q.push(i);

とする

RMQの値にインデックスを持たせる

$N$ 要素の数列

$A_0,A_1,\dots ,A_{N-1}$ があります。
以下のクエリに答えてください。

• $t=0$ のとき、
  $A_i$ に
  $x$ を代入
• $t=1$ のとき、
  $A[l,r)$ の最大値と、最大値をとる最小のインデックスを出力

Before

以下 https://ei1333.github.io/luzhiled/snippets/structure/segment-tree.html のセグメント木を使います。

int N;
cin >> N;
SegmentTree<pair<int, int>> rmq(N, [](auto a, auto b){ return max(a, b); }, {INT_MIN, 0});
for(int i = 0; i < N; i++){
    int A;
    cin >> A;
    rmq.set(i, {A, i});
}
rmq.build();
int Q;
cin >> Q;
for(int i = 0; i < Q; i++){
    int type, a, b;
    cin >> type >> a >> b;
    if(type == 0) rmq.update(a, {b, a});
    else{
        pair<int, int> ans = rmq.query(a, b);
        cout << ans.first << ' ' << ans.second << endl;
    }
}

After

int N;
cin >> N;
vector<int> A(N);
for(auto& a : A) cin >> a;
A.push_back(INT_MIN);
SegmentTree<int> rmq(N, [&](int a, int b){ return max(a, b, [&](int x, int y){ return A[x] < A[y]; }); }, N);
// max にも比較関数を入れられる　return A[a] < A[b] ? b : a; としても良い
for(int i = 0; i < N; i++) rmq.set(i, i); // 値の代わりにインデックス を持たせる
rmq.build();
int Q;
cin >> Q;
for(int i = 0; i < Q; i++){
    int type, a, b;
    cin >> type >> a >> b;
    if(type == 0){
        A[a] = b;
        rmq.update(a, a); // 代入し直して更新
    }
    else{
        int ans = rmq.query(a, b);
        cout << A[ans] << ' ' << ans << endl;
    }
}

グラフの辺にインデックスを持たせる

ここでは有向グラフとします。

Before

vector<vector<pair<int, int>>> g(N);
// として、 first を行先、 second を index とする
for(int i = 0; i < M; i++){
    int a, b;
    cin >> a >> b;
    g[a].push_back({b, i});
}

After

vector<int> to(M);
vector<vector<int>> g(N);
// として、辺のインデックスを入れる
for(int i = 0; i < M; i++){
    int a, b;
    cin >> a >> b;
    to[i] = b;
    g[a].push_back(i);
}

グラフに関しては struct Edge が強力なのであまり出番がなさそう

辺がたくさん属性を持っていて、それらを別々の配列で管理したいときに便利
メモリ効率が良い
pair を使わない
無向グラフの場合は有向グラフになおすとできる　状況によっては少し面倒になるかも

まとめ

• 値の代わりにインデックスを持たせることで、pair や tuple などを使うことなく実装ができる
• 考えることが減ってうれしい

テンプレート追加
vector<ll> iota(ll n){ vector<ll> ans(n); iota(ans.begin(), ans.end(), 0); return ans; }

— tatyam (@tatyam_prime) 2020年5月3日