せっかくQiitaのアカウントを取得しましたので、AtCoderの問題の解説を書いてみたいと思います。
問題
HHKB プログラミングコンテスト 2022（AtCoder Beginner Contest 235）
考察
やりたいこと
まずは落ち着いて問題文を読む。
長さNの数列 A がある。たとえば、A = [ 1 1 2 3 1 2 ] (入力例1の場合) 。
この数列に対して、以下のようなクエリが飛んでくるので、順に答える。
・「1が1番目に出てくる場所はどこか?」 → 「1番目です。」
・「2が2番目に出てくる場所はどこか?」 → 「6番目です。」
・「3が2番目に出てくる場所はどこか?」 → 「ありません。(3は1個しかない)」
・「5が1番目に出てくる場所はどこか?」 → 「ありません。(5は一つもない)」
愚直な探索の計算量
一つのクエリが来たら、Aの先頭から順に要素を見ていって、たとえば「3が2番目に出てくる場所」がどこかを調べれば、Aの要素へのアクセスを $2\times10^5$ 回程度することで、クエリに答えることが出来る。
しかし、クエリは $2\times10^5$ 個あるので、最悪で $(2\times10^5)\times(2\times10^5)=4\times10^{10}$ 回のAの要素へのアクセスが必要となる。
計算の効率化
上記の方法だと、毎回先頭から順に要素を見ている部分が効率が悪く感じられる。
極端な話、仮に最初に「1が3番目に出てくる場所はどこか」というクエリがあったなら、先頭から順に要素を調べて、3つめの「1」が見つかった時点で、「1が1番目に出てくる場所」「1が2番目に出てくる場所」は把握できている。
そもそも、先頭から順に調べる際に、「1」だけに注目するのではなく、他の数字もどこにあったかを覚えておくことも出来る。
それならいっそのこと、最初に先頭から最後までAを走査(スキャン)して、どの数字が、何番目と何番目にあったかのメモを作ればよい。
そうすると、クエリに対して(一つ一つAの要素を順に見ずに)即座に答えられる。
・「2が2番目に出てくる場所はどこか?」
→ 2は「3,6」番目にあるので、答えは「6番目」。
・「3が2番目に出てくる場所はどこか?」
→ 3は「4」番目にある。このリストは要素数が1なので、2番目の「3」は「ありません」。
・「5が1番目に出てくる場所はどこか?」
→ メモの中には5は登場していないので、「ありません」。
このようなデータ構造を作ることが出来れば、効率的に解ける。
必要な計算量は、
・メモの作成の際に、Aの要素に合計 $2\times10^5$ 回アクセスして内容を読み取る。
・$2\times10^5$ 個のクエリに対して、即座に答える。
なので、時間制限に十分間に合いそうだ。
実装
上記のようなメモは、「リストの辞書」を使うと実装できる。利用時に
・memo[1] → [1, 2, 5]
・memo[2] → [3,6]
・memo[3] → [4]
・memo[5] → (要素なし)
のように使えるようなイメージだ。
memoを作る手順は以下のようにする。先頭から順番にAの要素を見ていって、
・「1」番目は 1 だった → memo[1] のリストに「1」を追加。
・「2」番目は 1 だった → memo[1] のリストに「2」を追加。
・「3」番目は 2 だった → memo[2] のリストに「3」を追加。
・ (以下同様...)
のようにすればよい。
このような場合は、collections の defaultdict が役に立つ。
詳細は、下記のコードを見ていただければわかるかと思います。
解答
from collections import defaultdict
# 入力を受け取る
n,q = tuple(map(int, input().split()))
a = list(map(int, input().split()))
xk = [ tuple(map(int, input().split())) for _ in range(q)]
# メモの辞書
memo = defaultdict(list)
# メモを作る
for i,ai in enumerate(a):
memo[ai].append((i+1)) # python は 0-indexed なので、i に 1 を足しておく
# 答えを格納する配列
ans = []
for x,k in xk:
# A の中に x が存在する個数 = memo[x] の要素数 となる。
# A の中に x が存在しない場合、 memo[x] = [] (空リスト) になるので、問題ない。
if len(memo[x])<k:
ans.append(-1)
else:
ans.append(memo[x][k-1]) # python は 0-indexed なので、k から 1 を引く
# 答えを出力
print(*ans, sep='\n')