Official

A - お弁当の注文 / Lunch Box Orders Editorial by admin

Gemini 3.0 Flash (Thinking)

概要

\(N\) 種類のおかずの価格が与えられた状態で、指定された複数のおかずを組み合わせた「お弁当」の合計金額を \(M\) 件の注文ごとに計算する問題です。

考察

この問題の核となるのは、「各注文で指定されたおかずの価格をいかに効率よく取得し、合計するか」という点です。

1. データの管理

おかずの種類 \(N\) は最大 \(2 \times 10^5\)、おかずの価格 \(W_i\) も同様の数だけ存在します。これらを配列(Pythonではリスト)に格納することで、おかずの番号 \(i\) をインデックスとして、その価格 \(W_i\) を \(O(1)\) で取得できます。

2. 計算量の見積もり

各注文 \(j\) において、選ばれたおかずの数は \(K_j\) 個です。すべての注文に対する合計金額を求めるためには、延べ \(\sum K_j\) 回の加算処理が必要になります。 制約を見ると \(\sum K_j \leq 10^6\) となっており、これは単純な合計計算を行えば十分に実行時間制限内に間に合う量です。

3. 入出力の高速化

本問題では入力される数値の総数が非常に多く(最大で \(10^6\) 個以上)、Pythonの標準的な input() 関数を繰り返し呼び出すと、入出力処理だけで時間を消費し、実行時間制限(TLE)に達する可能性があります。 そのため、sys.stdin.read().split() を用いて入力を一括で読み込むなどの高速化手法が有効です。

アルゴリズム

  1. 価格表の作成: おかずの価格 \(W_1, W_2, \dots, W_N\) をリストに格納します。このとき、おかず番号が 1 から始まるため、リストのサイズを \(N+1\) にして 1-indexed で管理すると実装がスムーズです。
  2. 注文の処理:
    • 各注文 \(j\) について、指定されたおかずの番号 \(A_{j, 1}, A_{j, 2}, \dots, A_{j, K_j}\) を順に確認します。
    • 価格表からそれぞれの価格を引き出し、その合計値を計算します。
  3. 出力: すべての注文の計算結果を順番に出力します。

計算量

  • 時間計算量: \(O(N + \sum K_j)\)
    • 価格表の作成に \(O(N)\)、すべての注文のおかずを走査するのに \(O(\sum K_j)\) かかります。制約の範囲内であれば、およそ \(10^6\) 程度の操作となり、高速に動作します。
  • 空間計算量: \(O(N + \sum K_j)\)
    • おかずの価格と、入力されたすべての注文データを保持するためのメモリが必要です。

実装のポイント

  • 高速な入出力: sys.stdin.read().split() で全データを取得し、整数に変換して処理することで、大量の数値を効率的に扱えます。

  • リストのスライスと sum: Pythonでは、sum(W[id] for id in list_of_ids) のようにジェネレータ式を使ったり、スライスを工夫したりすることで、簡潔かつ高速に合計を求めることができます。

  • 一括出力: print を何度も呼び出すのではなく、結果をリストに蓄積して最後に '\n'.join() でまとめて出力することで、出力処理を高速化できます。

    ソースコード

import sys

def main():
    # 入力を一括で読み込み、空白で分割して整数のリストに変換します。
    # これにより、大量の入力を高速に処理することが可能です。
    try:
        input_str = sys.stdin.read().split()
        if not input_str:
            return
        data = list(map(int, input_str))
    except EOFError:
        return

    # N: おかずの種類数, M: 注文の件数
    N = data[0]
    M = data[1]

    # おかずの価格 W を 1-indexed で保持します。
    # data[2] から data[N+1] までがおかず 1 から N の価格に対応します。
    W = [0] * (N + 1)
    W[1:N+1] = data[2:N+2]

    # pos: 現在処理している注文の K_j (おかずの数) が格納されているインデックス
    pos = N + 2
    results = []

    # 各注文について合計金額を計算します。
    for _ in range(M):
        # K はこの注文で選ばれたおかずの数です。
        K = data[pos]
        
        # data[pos+1] から data[pos+K] までが選ばれたおかずの番号です。
        # スライスとジェネレータ式を用いて、合計金額を効率的に算出します。
        current_order_sum = sum(W[dish_id] for dish_id in data[pos + 1 : pos + 1 + K])
        
        # 計算結果を文字列としてリストに追加します。
        results.append(str(current_order_sum))
        
        # 次の注文の K_j の位置までポインタを進めます。
        pos += K + 1

    # すべての結果を改行で連結して一括出力します。
    sys.stdout.write('\n'.join(results) + '\n')

if __name__ == '__main__':
    main()

この解説は gemini-3-flash-thinking によって生成されました。

posted:
last update: