Official
D - お土産の組み合わせ / Souvenir Combinations Editorial by admin
Qwen3-Coder-480B概要
\(N\) 種類のお菓子と \(M\) 種類の飲み物からなるすべての組み合わせの「満足度」(お菓子の美味しさ × 飲み物の美味しさ)のうち、大きい方から \(K\) 個の合計を求めよという問題です。
考察
まず、単純に考えると、すべての組み合わせ \(N \times M\) 個の満足度をすべて計算して、大きい順にソートし、上位 \(K\) 個を選ぶ方法があります。しかし、制約より \(N \times M\) は最大で \(4 \times 10^{10}\) となり、これでは時間的に間に合いません(TLE)。
そこで効率的な方法を考えます。満足度は \(A_i \times B_j\) であり、\(A\) と \(B\) どちらも大きい値同士を掛けたものが最も大きな満足度になります。したがって、\(A\) と \(B\) をともに降順にソートしておくことで、最大の満足度は必ず \(A_0 \times B_0\) になります。
このように「順序がある」という性質を利用して、全探索せずに上位の組み合わせを効率よく取り出す必要があります。ここでは、優先度付きキュー(ヒープ) を使って、現在の最大値から次に大きな候補を探索的に取り出していきます。
具体的には、組み合わせを \((i, j)\) のペアで管理し、\((0, 0)\) から初めて、そこから次に候補となり得る \((i+1, j)\) と \((i, j+1)\) をヒープに追加していきます。ただし、同じ組み合わせを二度追加しないよう、使用済みチェックを行います。
このようにして、上位 \(K\) 個の満足度を効率的に取り出し、合計を求めることができます。
アルゴリズム
- お菓子の美味しさリスト \(A\) と飲み物の美味しさリスト \(B\) をともに降順にソートします。
- 優先度付きキュー(最大ヒープを実現するために負の値を使う)を用意し、最初の最大値となる組み合わせ \((0, 0)\) を追加します。
- 同時に、追加済みの組み合わせを記録する
setを用意します。 - \(K\) 回だけ以下を繰り返します:
- ヒープから最大の満足度の組み合わせを取り出し、その値を合計に加える。
- 取り出した組み合わせの「右 \((i, j+1)\)」と「下 \((i+1, j)\)」の候補を新たにヒープに追加(すでに追加されていない場合のみ)。
- 合計値を出力します。
計算量
- 時間計算量: \(O(K \log K)\)
(ヒープへの挿入・取り出しは \(O(\log K)\) で、これを最大 \(K\) 回行う) - 空間計算量: \(O(K)\)
(ヒープおよび使用済みチェック用の集合のサイズは最大 \(K\))
実装のポイント
- Pythonの
heapqは最小ヒープなので、最大値を取り出すために値を負にして格納する。 - 同じ \((i, j)\) の組み合わせを複数回ヒープに追加しないために
setを使う。 - 最初に \(A\) と \(B\) を降順にソートすることで、最大値から探索を始められる。
## ソースコード
```python
import heapq
def main():
import sys
input = sys.stdin.read
data = input().split()
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, B を降順にソート
A.sort(reverse=True)
B.sort(reverse=True)
# 使用済みチェック用のセット (i, j) の形式
used = set()
# 最大ヒープ(Pythonは最小ヒープなので負の値を入れる)
heap = []
# 最大の組み合わせを最初に追加
heapq.heappush(heap, (-A[0] * B[0], 0, 0))
used.add((0, 0))
total = 0
for _ in range(K):
if not heap:
break
neg_val, i, j = heapq.heappop(heap)
val = -neg_val
total += val
# 次の候補を追加(隣接するインデックス)
if i + 1 < N and (i + 1, j) not in used:
heapq.heappush(heap, (-A[i + 1] * B[j], i + 1, j))
used.add((i + 1, j))
if j + 1 < M and (i, j + 1) not in used:
heapq.heappush(heap, (-A[i] * B[j + 1], i, j + 1))
used.add((i, j + 1))
print(total)
if __name__ == "__main__":
main()
この解説は qwen3-coder-480b によって生成されました。