概要

各注文について、選ばれたおかずの番号に対応する価格を合計して出力する、シンプルな集計問題です。

考察

この問題は一見すると非常にシンプルですが、データサイズに注意が必要です。

• おかずは最大 \(N = 2 \times 10^5\) 種類
• 注文は最大 \(M = 10^5\) 件
• 全注文のおかず数の合計 \(\sum K_j\) は最大 \(10^6\)

素朴に各注文ごとにおかずの価格を足し合わせるだけで十分ですが、重要な気づきは 全注文を通じたおかず数の合計が \(10^6\) 以下 という制約です。これにより、注文ごとに \(K_j\) 個のおかずを走査しても、全体の処理回数は \(10^6\) 回以下に収まります。

例えば、以下のような入力を考えます：

• おかずが 3 種類で、価格が \(W = [100, 200, 300]\)
• 注文1：おかず 1, 3 → 合計 \(100 + 300 = 400\)
• 注文2：おかず 2 → 合計 \(200\)

各注文で指定されたおかず番号をインデックスとして価格配列 \(W\) を参照し、合計するだけです。

アルゴリズム

1. \(N\) 種類のおかずの価格を配列 \(W\) に格納する（0-indexed）。
2. \(M\) 件の注文それぞれについて以下を行う：
   • 注文に含まれる \(K_j\) 個のおかず番号 \(A_{j,1}, A_{j,2}, \dots, A_{j,K_j}\) を読み取る。
   • 各おかず番号 \(A_{j,k}\) に対して \(W[A_{j,k} - 1]\)（1-indexed → 0-indexed に変換）を合計する。
   • 合計金額を出力用に保存する。
3. すべての注文の合計金額をまとめて出力する。

計算量

• 時間計算量: \(O\!\left(N + \displaystyle\sum_{j=1}^{M} K_j\right)\)
  • 価格配列の読み込みに \(O(N)\)、全注文の処理に合計 \(O\!\left(\sum K_j\right)\)（最大 \(10^6\)）
• 空間計算量: \(O(N + M)\)
  • 価格配列 \(W\) に \(O(N)\)、出力用配列に \(O(M)\)

実装のポイント

• 入力の高速化: Pythonでは標準の input() は遅いため、sys.stdin.readline に置き換えることで高速化しています。大量の入力がある問題では必須のテクニックです。

• 出力の高速化: 各注文ごとに print() を呼ぶと \(M\) 回の出力で遅くなるため、結果を配列に溜めて最後に sys.stdout.write() で一括出力しています。

• インデックスの変換: おかず番号は 1-indexed で与えられますが、Pythonのリストは 0-indexed なので、W[line[i] - 1] のように \(-1\) して参照する必要があります。

• 入力行のパース: 各注文の行は先頭が \(K_j\)（おかず数）で、続く \(K_j\) 個がおかず番号です。一度にリストとして読み取り、line[0] を \(K_j\)、line[1] 〜 line[K_j] をおかず番号として処理しています。

  ソースコード

import sys
input = sys.stdin.readline

def main():
    N, M = map(int, input().split())
    W = list(map(int, input().split()))
    out = []
    for _ in range(M):
        line = list(map(int, input().split()))
        K = line[0]
        total = 0
        for i in range(1, K + 1):
            total += W[line[i] - 1]
        out.append(str(total))
    sys.stdout.write('\n'.join(out) + '\n')

main()

この解説は claude4.6opus-thinking によって生成されました。