問題文
解説
標準入力
本問題の入力される値は次の形式で与えられます。
N M
S_1 P_1
S_2 P_2
.
.
.
S_N P_N
Q_1
Q_2
.
.
.
Q_M
入力の受け取り方については過去に次のような記事を書いているので参考にしてみてください。
この記事によると次のように入力を受け取れば良いことがわかります。
# 標準入力
n, m = map(int, input().split())
sp = [list(input().split()) for _ in range(n)]
q = [input() for _ in range(m)]
計算量
計算量の観点から問題を捉えるためには制約に着目します。
本問題の制約は次のようになります。
$1 \le N \le 10^4$
$1 \le M \le 10^4$
$1 \le |S_i| \le 10^2$
$1 \le P_i \le 10^4$
$1 \le |Q_i| \le 10^2$
アルゴリズム 1
まずは、一つ目のクエリ $Q_1$ に対してどのように処理を行えば良いか考えることにします。
文字列 $S_i$ のうち、先頭が $Q_1$ で始まる文字列を選ぶ必要があるので、 $S_i$ の $j$ 文字目と $Q_1$ の $j$ 文字目が等しいかどうかを判定するアルゴリズムが考えられますね。
$Q_1$ の $1$ 文字目から $|Q_1|\ (\le 10^2)$ 文字目までを比較するので、一つ目のクエリを処理するのに要する計算量は $O(N|Q_1|)$ となります。こういったクエリが $M$ 個あるので、全体で $O(NM\max(|Q_i|))$ となります。
最悪のケースを考えると $N = 10^4, M = 10^4, \max(|Q_i|) = 10^2$ となるので $O(NM\max(|Q_i|)) = 10^4 \times 10^4 \times 10^2 = 10^{10}$ となります。
個人的には、 $10^7$ より大きい場合はそのアルゴリズムで合格するのは厳しいと感じています。ひょっとしたら $10^8$ 以下であれば合格できる可能性もありますが、 Paiza の仕様には詳しくないのでここでは $10^7$ を上限とします。
アルゴリズム 2
計算量が $10^7$ 以下となるようにアルゴリズムを改善してみましょう。
本問題の制約を再掲します。
$1 \le N \le 10^4$
$1 \le M \le 10^4$
$1 \le |S_i| \le 10^2$
$1 \le P_i \le 10^4$
$1 \le |Q_i| \le 10^2$
$M$ 個のクエリは当然ながら全て処理する必要があるので、どのような工夫を行っても計算量は $O(M)$ 以上となりますね。 $O(M) = 10^4$ なので、 $O(NM)$ のように上限値が $10^4$ であるようなもの（例えば $N$ ）を $M$ と掛け算すると $10^4 \times 10^4 = 10^8$ となり、 $10^7$ 以下に抑えることができません。
したがって、上限値が $10^2$ であるようなものを $1$ つだけ用いるようなアルゴリズムを設計すれば、 $10^4 \times 10^2 = 10^6$ となり、 $10^7$ 以下に抑えることができそうです。
まず、文字列 $S_1$ について考えます。
先頭がどのような文字列ならば、文字列 $S_1$ はその対象となるか考えると、 $1 \le k \le |S_1|$ に対して、 $S_1$ の先頭から $k$ 文字分切り取った文字列であることが分かるかと思います。
文字列 $S_1$ の先頭から $k$ 文字分切り取った文字列を $S_1[:k]$ とします。
このとき、文字列 $S_1$ に対して、 $S_1[:1], S_1[:2], \dots, S_1[:|S_1|]$ を求めるのに $O(|S_1|)$ あれば十分ですね。
この部分は少し抽象的でイメージしにくいかもしれません。具体例を挙げます。
例えば、 $S_1 =$ qiita であるとき、 $S_1[:1], S_1[:2], \dots, S_1[:|S_1|]$ はそれぞれ、 q, qi, qii, qiit, qiita を表しています。
文字列 $S_i$ に対しても同様に考えると、先頭の候補を全て列挙するのに必要な計算量は $O(N\max(|S_i|))$ となるので、 $N = 10^4, \max(|S_i|) = 10^2$ より、 $O(N\max(|S_i|)) = 10^4 \times 10^2 = 10^6 \le 10^7$ となりますね。
列挙した全ての候補に対して価格の情報を付け加えれば、各クエリに対して $O(1)$ で必要な金額を求めることができます。これが実現可能なデータ構造として辞書（連想配列）があります。
辞書に対する主な操作の計算量は下表のようになります。こちらのサイト の内容をまとめたものになります。
| 操作 | （平均）計算量 |
|---|---|
| アイテムの取得 | $O(1)$ |
| アイテムの追加 | $O(1)$ |
| アイテムの削除 | $O(1)$ |
以上のことから、この問題を解くのに必要な計算量は $O(N\max(|S_i|) + M) \le 10^7$ となるので合格となります。
提出コード（Python）
dict （辞書）の代わりに defaultdict を使うと実装が少し楽になります（キーの存在チェックが不要になるため）。
from collections import defaultdict
# 標準入力
n, m = map(int, input().split())
sp = [list(input().split()) for _ in range(n)]
q = [input() for _ in range(m)]
# アルゴリズム 2
d = defaultdict(int) # 辞書
for i in range(n):
s, p = sp[i][0], int(sp[i][1])
substr = "" # 先頭から k 文字分切り取った文字列を格納する
for si in s:
substr += si
d[substr] += p # 先頭が substr で始まる文字列すべてを買い占めるのに必要な金額を辞書で管理する
for qi in q:
ans = d[qi] # 先頭が qi で始まる文字列すべてを買い占めるのに必要な金額は d[qi] で確認できる
print(ans)
余談
ご飯を食べながらスマホで問題文を読み、計算量を意識しながらどういったアルゴリズムであれば通用するか $10$ 分ほど考え、考えたものを食後にそのまま実装してみたら合格できました。このように計算量を考えると、コードを書く前から適切なアルゴリズムを設計したり選定したりするのに役立つことが分かるかと思います。本記事は食事中に考えていたことを言語化したものになります。是非、計算量を意識しましょう！