C - The Kth Time Query - Harryのブログ

皆さん、こんにちは！

ハリーです。今回は、C - The Kth Time Query を解いていきたいと思います。この問題は HHKB プログラミングコンテスト 2022（AtCoder Beginner Contest 235）で出題された問題です。私もこのコンテストに参加していたのですが、解くことができなかったので復習していきたいと思います！

C - The Kth Time Query
実装方法
ACコード
おわりに

C - The Kth Time Query

問題文

長さ $N$ の数列 $A = (a_1,a_2,...,a_N)$ があります。
以下で説明される $Q$ 個のクエリに答えてください。

クエリ $i :$ 整数の組 $(x_i ,k_i)$ が与えられます。$A$ の要素を $a_1 ,a_2 ,…$ と前から順に見ていったときに、数 $x_i$が $k_i$ 回目に登場するのは $A$ の前から何番目の要素を見たときかを出力してください。
ただし条件を満たす要素が存在しない場合は $−1$ を出力してください。

制約

$1 \leq N \leq 2 × 10^5$
$1 \leq Q \leq 2 × 10^5$
$0 \leq a_i \leq 10^9(1 \leq i \leq N)$
$0 \leq x_i \leq 10^9(1 \leq i \leq Q)$
$1 \leq k_i \leq N(1 \leq i \leq Q)$
入力はすべて整数である。

入力例１

6 8
1 1 2 3 1 2
1 1
1 2
1 3
1 4
2 1
2 2
2 3
4 1

出力例１

入力例２

3 2
0 1000000000 999999999
1000000000 1
123456789 1

出力例２

2
-1

実装方法

まず初めに、collections ライブラリのimport を行います。

 from collections import defaultdict

次に入力を受け取ります。

 N, Q = map(int, input().split())
 a = list(map(int, input().split()))

ここで、defaultdict を作成して引数に d という関数を設定します。

(※defaultdict の参考記事はこちらです）

 d = defaultdict(list)

配列 d の受け取った a の要素の index 番目に値を加えていきます。

 for i in range(N):
   d[a[i]].append(i + 1)

$x$ と $k$ を受け取り、d[x] の長さ（length) が $k$ 以上の場合は、d[x] の前から $k$ 番目の値を出力します。d[x] の長さが $k$ 未満の場合は、$-1$ を出力します。この工程を $Q$ 回繰り返します。

 for _ in range(Q):
   x, k = map(int, input().split())
   if len(d[x]) >= k:
       print(d[x][k - 1])
     else:
       print(-1)

ACコード

 from collections import defaultdict

 N, Q = map(int, input().split())
 a = list(map(int, input().split()))

 d = defaultdict(list)

 for i in range(N):
   d[a[i]].append(i + 1)

 for _ in range(Q):
   x, k = map(int, input().split())
   if k <= len(d[x]):
       print(d[x][k - 1])
   else:
       print(-1)

おわりに

どうでしたでしょうか？今回は defaultdict を使う問題を解いてみました。defaultdict を使った問題はコンテストでもよく出題されるので、しっかり復習することができて良かったです！次の記事でも、もう1問くらい解いていきたいと思います。