Official

B - 商品選抜 / Product Selection Editorial by admin

Qwen3-Coder-480B

概要

立ち入り可能なエリアに保管されている商品の中から、価格の高いものから最大 \(K\) 個選び、その合計価格の最大値を求めます。

考察

この問題では、まず「どの商品を選ぶことができるか」が重要です。
商品はそれぞれエリアに属しており、その中でも立ち入り可能なエリアに属する商品だけが選択可能です。

したがって、以下のステップが必要になります:

  1. 立ち入り可能なエリアを記録する(集合や配列など)。
  2. 各商品について、それが立ち入り可能なエリアにあるかどうかを判定。
  3. 条件を満たす商品の価格を集め、その中から最も高い価格のものを \(K\) 個選ぶ。

素朴な方法としては、すべての商品を見て、それぞれのエリアが許可されているか調べるというやり方がありますが、これは十分高速です(後述)。
しかし、もし価格のソートや選別を適切に行わないと、誤った結果や無駄な計算が増えてしまいます。

例えば、価格を降順にソートせずに上位 \(K\) 個を選ぶと、本来より低い合計額が出てしまう可能性があります。

アルゴリズム

  1. 入力を読み込み、各商品の価格 \(V_i\) と所属エリア \(C_i\) を保存します。
  2. 立ち入り可能なエリア番号 \(S_j\) をセット型(set)などで管理して、高速な検索を可能にします。
  3. すべての商品について、その所属エリア \(C_i\) が立ち入り可能であるかをチェックします。
  4. 立ち入り可能なエリアにある商品の価格のみをリストに追加します。
  5. このリストを降順にソートし、先頭から \(K\) 個の価格を取り出して合計します。

具体例

例えば以下のような入力があったとします:

N=5, M=3, K=3, Q=2
価格: [100, 200, 50, 300, 150]
エリア: [1, 2, 3, 1, 2]
立ち入り可能: [1, 2]

このとき、エリア 1 または 2 にある商品の価格は [100, 200, 300, 150] の4つ。
これらを大きい順にソートすると [300, 200, 150, 100]。
上位3つを取ると合計は \(300 + 200 + 150 = 650\) となり、これが答えになります。

計算量

  • 時間計算量: \(O(N \log N)\) (価格リストのソートが支配的)
  • 空間計算量: \(O(N + Q)\) (商品リストと立ち入り可能エリアの保存)

実装のポイント

  • エリア番号の存在確認には set を使うことで、\(O(1)\) での判定が可能。

  • 入力が多いので sys.stdin.read を使用して高速に読み込むと良い。

  • ソート時に reverse=True を指定して降順ソートすることに注意。

    ソースコード

import heapq
import sys

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    idx = 0
    N = int(data[idx]); idx += 1
    M = int(data[idx]); idx += 1
    K = int(data[idx]); idx += 1
    Q = int(data[idx]); idx += 1
    
    V = [0] * N
    C = [0] * N
    for i in range(N):
        V[i] = int(data[idx]); idx += 1
        C[i] = int(data[idx]); idx += 1
    
    S = set()
    for _ in range(Q):
        s = int(data[idx]); idx += 1
        S.add(s)
    
    # 立ち入り可能なエリアにある商品の価格をリスト化
    available_prices = []
    for i in range(N):
        if C[i] in S:
            available_prices.append(V[i])
    
    # 価格が高い順にソート
    available_prices.sort(reverse=True)
    
    # 上位K個の合計を計算
    max_sum = sum(available_prices[:K])
    
    print(max_sum)

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。