A - お弁当の注文 / Lunch Box Orders Editorial by admin
GPT 5.4 High概要
各注文について、「選ばれたおかずの番号」を見ながら対応する価格を足し合わせるだけです。
制約を見ると、全注文に含まれるおかず数の合計が \(\sum K_j \le 10^6\) なので、各指定をそのまま順に処理すれば十分間に合います。
考察
この問題で欲しいのは、各注文に対して
- 選ばれたおかず \(A_{j,1}, A_{j,2}, \dots, A_{j,K_j}\)
- それぞれの価格 \(W_{A_{j,1}}, W_{A_{j,2}}, \dots, W_{A_{j,K_j}}\)
を足した合計です。
つまり、\(j\) 件目の答えは
\(W_{A_{j,1}} + W_{A_{j,2}} + \cdots + W_{A_{j,K_j}}\)
になります。
重要な気づき
おかずの価格 \(W_i\) は最初に全部与えられているので、
「おかず番号 \(\to\) 価格」の対応を配列に持っておけば、各注文のおかず番号を読むたびにすぐ価格を取り出せます。
たとえば、
- 価格: \(W_1=300, W_2=500, W_3=200\)
- 注文: \([1, 3]\)
なら、合計は \(300 + 200 = 500\) です。
素朴なまずい方法
各注文ごとに「全 \(N\) 種類のおかずを毎回確認して、そのおかずが注文に含まれるか」を調べる方法だと、
1 件あたり \(O(N)\) 以上かかり、全体では \(O(NM)\) になります。
制約は
- \(N \le 2 \times 10^5\)
- \(M \le 10^5\)
なので、\(O(NM)\) は大きすぎて間に合いません。
どう解決するか
注文に実際に含まれているおかずだけを見ればよいです。
各注文について、列挙された \(K_j\) 個のおかず番号を順に読み、その価格を足し合わせます。
このとき全体で処理するおかず番号の個数は \(\sum K_j \le 10^6\) なので、
全体計算量は \(O(\sum K_j)\) となり十分高速です。
アルゴリズム
- おかずの価格を配列
Wに保存する。 - 各注文について:
- 個数 \(K\) を読む
- 合計
s = 0を用意する - 続く \(K\) 個のおかず番号
aについて、s += W[a]とする - 最後に
sを答えとして保存する
- すべての答えを改行区切りで出力する
この方法なら、各おかず番号に対して 1 回ずつ価格を足すだけです。
計算量
- 時間計算量: \(O\left(N + \sum_{j=1}^{M} K_j\right)\)
- 空間計算量: \(O(N + M)\)
実装のポイント
入力サイズが大きいので、
sys.stdin.buffer.read().split()を使ってまとめて高速に読み込んでいます。おかず番号は \(1\) から始まるので、コードでは
W = [0] + ...として 1-indexed にしています。
これにより、おかず番号iに対してそのままW[i]で価格を取り出せます。出力も 1 行ずつ
printする代わりに、文字列配列にためて最後に"\n".join(ans)でまとめて出力すると高速です。ソースコード
import sys
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
it = iter(data)
N = next(it)
M = next(it)
W = [0] + [next(it) for _ in range(N)]
ans = []
for _ in range(M):
K = next(it)
s = 0
for _ in range(K):
s += W[next(it)]
ans.append(str(s))
sys.stdout.write("\n".join(ans))
if __name__ == "__main__":
main()
この解説は gpt-5.4-high によって生成されました。
posted:
last update: