Overview

Given \(N\) types of snacks and \(M\) types of drinks, you choose one of each to form a combination with a satisfaction score (product of deliciousness values). The task is to find the sum of the top \(K\) largest satisfaction scores.

Analysis

Naive Approach and Its Issues

The simplest method is to compute all \(N \times M\) products, sort them, and take the top \(K\). However, since \(N\) and \(M\) can each be up to \(2 \times 10^5\), the number of combinations can reach \(4 \times 10^{10}\), which is far too large to handle.

Key Insight

If we sort \(A\) and \(B\) in descending order, the following properties hold:

• \(A[0] \times B[0]\) is the maximum among all combinations
• For \(A[i] \times B[j]\), the next largest candidates are either \(A[i+1] \times B[j]\) or \(A[i] \times B[j+1]\)

For example, if \(A = [5, 3, 1]\), \(B = [4, 2]\):

	\(B[0]=4\)	\(B[1]=2\)
\(A[0]=5\)	20	10
\(A[1]=3\)	12	6
\(A[2]=1\)	4	2

After taking the maximum \(20\) (\(i=0, j=0\)), the candidates are \(12\) (\(i=1, j=0\)) and \(10\) (\(i=0, j=1\)). In this way, by only adding the elements “to the right” and “below” the currently extracted element as candidates, we can efficiently extract the top \(K\) elements without exhaustive search.

Solution Method

This is a classic problem of “extracting the top \(K\) elements from a sorted matrix,” which can be solved using a priority queue (heap).

Algorithm

1. Sort \(A\) and \(B\) each in descending order.
2. Insert the initial state \((A[0] \times B[0],\ 0,\ 0)\) into a max-heap (in Python, use a min-heap with negated values).
3. Repeat the following \(K\) times:
   • Extract the \((i, j)\) with the maximum value from the heap, and add \(A[i] \times B[j]\) to the running total.
   • If \((i+1, j)\) is unvisited and within bounds, add it to the heap and mark it as visited.
   • If \((i, j+1)\) is unvisited and within bounds, add it to the heap and mark it as visited.
4. Output the total.

Preventing duplicates is important. For example, \((1, 1)\) can be reached both from “below” \((0, 1)\) and from “right” of \((1, 0)\), so we use a visited set to ensure the same \((i, j)\) is not added twice.

Complexity

• Time complexity: \(O(N \log N + M \log M + K \log K)\)
  • Sorting takes \(O(N \log N + M \log M)\).
  • We perform heap operations \(K\) times, adding at most 2 elements per iteration, so the heap size is at most \(O(K)\). Each operation is \(O(\log K)\).
• Space complexity: \(O(N + M + K)\)
  • Arrays \(A, B\) take \(O(N + M)\), and the heap and visited set take \(O(K)\).

Implementation Notes

• Python’s heapq is a min-heap, so we store negated values to use it as a max-heap.

• The visited set prevents duplicate additions of the same \((i, j)\). Without this, the same combination would be counted multiple times.

• Since the constraint is \(K \leq 2 \times 10^5\), the heap size does not explode and the algorithm runs efficiently.

• Fast input using sys.stdin.buffer.read() is employed to avoid TLE on large inputs.

  Source Code

import heapq
import sys

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    M = int(input_data[idx]); idx += 1
    K = int(input_data[idx]); idx += 1
    A = [int(input_data[idx + i]) for i in range(N)]; idx += N
    B = [int(input_data[idx + i]) for i in range(M)]; idx += M
    
    # Sort both in descending order
    A.sort(reverse=True)
    B.sort(reverse=True)
    
    # Use a max-heap (negate values for min-heap)
    # State: (-A[i]*B[j], i, j)
    # Start with (A[0]*B[0]) which is the maximum
    
    heap = [(-A[0] * B[0], 0, 0)]
    visited = set()
    visited.add((0, 0))
    
    total = 0
    for _ in range(K):
        neg_val, i, j = heapq.heappop(heap)
        total += -neg_val
        
        # Push (i+1, j) and (i, j+1)
        if i + 1 < N and (i + 1, j) not in visited:
            visited.add((i + 1, j))
            heapq.heappush(heap, (-A[i + 1] * B[j], i + 1, j))
        if j + 1 < M and (i, j + 1) not in visited:
            visited.add((i, j + 1))
            heapq.heappush(heap, (-A[i] * B[j + 1], i, j + 1))
    
    print(total)

main()

---

This editorial was generated by claude4.6opus-thinking.