概要

お菓子と飲み物の全ての組み合わせの満足度（美味しさの積）のうち、大きい方からK個の合計値を求める問題です。

考察

素朴なアプローチとして、全ての組み合わせ（\(N \times M\)通り）の満足度を計算し、ソートして上位K個を選ぶ方法が考えられます。しかし、\(N\)と\(M\)が最大で\(2 \times 10^5\)まで大きくなるため、組み合わせの数は最大\(4 \times 10^{10}\)通りとなり、計算量が膨大で現実的ではありません。

そこで、満足度が大きい組み合わせから効率的に探索する方法が必要です。お菓子と飲み物をそれぞれ美味しさの降順にソートすることで、満足度が大きい組み合わせから順に候補を生成できます。具体的には、各飲み物に対して最も美味しいお菓子との組み合わせから始め、次に美味しいお菓子との組み合わせを優先度付きキューで管理することで、効率的に上位K個を取得できます。

アルゴリズム

1. お菓子のリスト\(A\)と飲み物のリスト\(B\)をそれぞれ降順にソートします。
2. 優先度付きキュー（最小ヒープ）を初期化します。最初に、各飲み物\(B_j\)に対して、最も美味しいお菓子\(A_0\)との組み合わせの満足度（\(A_0 \times B_j\)）を、負の値（最小ヒープを最大ヒープのように使うため）とインデックス情報とともにキューに追加します。
3. キューから満足度が最大の組み合わせをポップし、合計値に加算します。
4. ポップした組み合わせのお菓子のインデックスが\(i\)、飲み物のインデックスが\(j\)の場合、次のインデックス\(i+1\)のお菓子と同じ飲み物\(j\)の組み合わせ（\(A_{i+1} \times B_j\)）をキューに追加します。
5. 上記の操作をK回繰り返し、合計値を出力します。

この手法では、各飲み物に対してお菓子を美味しい順に探索することで、全ての組み合わせを生成することなく上位K個を効率的に取得できます。

計算量

• 時間計算量: \(O(N \log N + M \log M + K \log M)\)
  • ソートに\(O(N \log N + M \log M)\)
  • 優先度付きキューの操作に\(O(K \log M)\)（キューサイズは最大\(M\)）
• 空間計算量: \(O(N + M)\)

実装のポイント

• 優先度付きキューを最小ヒープとして使い、負の値で格納することで最大値を効率的に取得します。

• 各組み合わせと一緒にインデックス情報を保存し、次の候補を生成できるようにします。

• キューには各飲み物に対して現在の最適なお菓子のインデックスを管理し、次のインデックスが存在する場合のみ次の候補を追加します。

  ソースコード

import sys
import heapq

def main():
    data = sys.stdin.read().split()
    if not data:
        return
    
    n = int(data[0])
    m = int(data[1])
    K = int(data[2])
    A = list(map(int, data[3:3+n]))
    B = list(map(int, data[3+n:3+n+m]))
    
    A.sort(reverse=True)
    B.sort(reverse=True)
    
    heap = []
    for j in range(m):
        heapq.heappush(heap, (-A[0] * B[j], 0, j))
    
    total = 0
    count = 0
    while count < K and heap:
        val, i, j = heapq.heappop(heap)
        total += -val
        count += 1
        if i + 1 < n:
            heapq.heappush(heap, (-A[i+1] * B[j], i+1, j))
    
    print(total)

if __name__ == "__main__":
    main()

この解説は deepseekv3 によって生成されました。