D - お土産の組み合わせ / Souvenir Combinations 解説 by admin
Qwen3-Coder-480BOverview
Given \(N\) types of snacks and \(M\) types of drinks, we need to find the sum of the top \(K\) “satisfaction values” (snack deliciousness × drink deliciousness) among all possible combinations.
Approach
First, a naive approach would be to compute all \(N \times M\) satisfaction values, sort them in descending order, and select the top \(K\). However, since \(N \times M\) can be up to \(4 \times 10^{10}\), this would not finish in time (TLE).
So we need a more efficient method. Since the satisfaction value is \(A_i \times B_j\), the largest satisfaction comes from multiplying the largest values of both \(A\) and \(B\). Therefore, by sorting both \(A\) and \(B\) in descending order, the maximum satisfaction is always \(A_0 \times B_0\).
By leveraging this “ordering” property, we need to efficiently extract the top combinations without exhaustive search. Here, we use a priority queue (heap) to exploratively extract the next largest candidates starting from the current maximum.
Specifically, we manage combinations as pairs \((i, j)\), starting from \((0, 0)\), and add the next potential candidates \((i+1, j)\) and \((i, j+1)\) to the heap. To avoid adding the same combination twice, we perform a visited check.
In this way, we can efficiently extract the top \(K\) satisfaction values and compute their sum.
Algorithm
- Sort both the snack deliciousness list \(A\) and the drink deliciousness list \(B\) in descending order.
- Prepare a priority queue (using negative values to simulate a max-heap), and add the initial maximum combination \((0, 0)\).
- Also prepare a
setto record which combinations have already been added. - Repeat the following \(K\) times:
- Extract the combination with the maximum satisfaction from the heap and add its value to the total.
- Add the “right \((i, j+1)\)” and “down \((i+1, j)\)” candidates of the extracted combination to the heap (only if they haven’t been added already).
- Output the total.
Complexity
- Time complexity: \(O(K \log K)\)
(Each insertion/extraction from the heap is \(O(\log K)\), and this is done at most \(K\) times) - Space complexity: \(O(K)\)
(The size of the heap and the visited set is at most \(K\))
Implementation Notes
Python’s
heapqis a min-heap, so we store negative values to extract the maximum.Use a
setto prevent adding the same \((i, j)\) combination to the heap multiple times.By sorting \(A\) and \(B\) in descending order first, we can begin the search from the maximum value.
Source Code
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()
This editorial was generated by qwen3-coder-480b.
投稿日時:
最終更新: