けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 235 C - The Kth Time Query (4Q, 灰色, 300 点)

連想配列の練習問題!

問題概要

長さ  N の整数列  a_{1}, a_{2}, \dots, a_{N} が与えられる。この数列に対して、次の  Q 個のクエリに答えよ。

【クエリ】
整数  x, k が与えられるので、整数列を順に見ていったときの「 k 個目の値  x の要素」の index を答えよ(存在しない場合は -1)

制約

  •  1 \le N, Q \le 2 \times 10^{5}
  •  0 \le a_{i}, x_{i} \le 10^{9}

考えたこと

数列の情報をそのまま扱うのではなく、クエリに答えやすい形になるように「前処理」しておきたい。具体的には、次のような配列を持ちたい。


  • ids[x]:数列中の、値が  x であるような index を小さい順に格納した集合(C++ ならば配列の値は vector<int> 型でよさそう)

しかしながら、この配列の添字は 0 以上  10^{9} 以下の数が入ることが考えられる。これはメモリ制限に引っかかる。そこで活躍するのが連想配列だ。

C++ であれば、map<(添字の型), (配列の値の型)> という型が使える。今回は map<int, vector<int>> 型がピッタリだ。なお、この map の各要素にアクセスするのに要する計算量は  O(\log N) である。よって、

  • 前処理に要する計算量: O(N \log N)
  • クエリの処理に要する計算量: O(Q \log N)

となる。全体の計算量は  O((N + Q) \log N) と評価できる。

コード

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

int main() {
    int N, Q, a, x, k;
    cin >> N >> Q;
    map<int, vector<int>> ids;
    for (int i = 0; i < N; i++) {
        cin >> a;
        ids[a].push_back(i+1);
    }

    while (Q--) {
        cin >> x >> k;
        if (ids[x].size() < k) cout << -1 << endl;
        else cout << ids[x][k-1] << endl;
    }
}