Problem Description

You have two arrays nums1 and nums2, both with the same length n. Your task is to select exactly k indices from these arrays to maximize a specific score.

The score is calculated as follows:

• Take the sum of all selected elements from nums1
• Multiply this sum by the minimum value among the corresponding selected elements from nums2
• Formula: (nums1[i₀] + nums1[i₁] + ... + nums1[iₖ₋₁]) × min(nums2[i₀], nums2[i₁], ..., nums2[iₖ₋₁])

For example, if you select indices 0, 2, and 4:

• Sum from nums1: nums1[0] + nums1[2] + nums1[4]
• Minimum from nums2: min(nums2[0], nums2[2], nums2[4])
• Score: sum × minimum

The selected indices form a subsequence, meaning you can pick any k indices from the arrays (not necessarily consecutive).

Your goal is to find the selection of k indices that gives the maximum possible score.

The solution approach sorts the pairs (nums2[i], nums1[i]) in descending order by nums2 values. Then it traverses through these sorted pairs, maintaining a min heap of size k containing elements from nums1. For each position, when the heap reaches size k, it calculates the score using the current sum and the current nums2 value (which acts as the minimum since we're traversing in descending order). The heap helps efficiently manage which k elements from nums1 to keep for the maximum sum at each step.

Intuition

The key insight is that once we fix the minimum value from nums2, we want to maximize the sum from nums1. This suggests we should consider each possible minimum value and find the best sum for that minimum.

Think about it this way: if we've already decided that element nums2[i] will be our minimum, then we know all selected elements from nums2 must be greater than or equal to nums2[i]. This means we can only choose indices where nums2[j] >= nums2[i].

This leads us to a greedy strategy: sort the pairs by nums2 values in descending order. As we traverse from left to right, each nums2 value we encounter could potentially be the minimum in our final selection. At each position i, all elements to the left have nums2 values greater than or equal to nums2[i], making them valid candidates to combine with.

Now the problem becomes: given that nums2[i] is our minimum, which k-1 elements should we pick from the left (plus the current element) to maximize the sum from nums1? We want the largest k elements from nums1 among all valid positions.

A min heap of size k perfectly solves this subproblem. As we traverse, we:

1. Add the current nums1 value to our heap
2. If heap size exceeds k, remove the smallest element
3. When heap has exactly k elements, calculate the score using the current nums2 value as the minimum and the heap sum as our nums1 sum

This way, at each position, we maintain the best possible k elements from nums1 that can be paired with the current minimum from nums2, ensuring we check all possible minimums and find the optimal score.

Learn more about Greedy, Sorting and Heap (Priority Queue) patterns.

Solution Approach

The implementation follows these steps:

1. Create and Sort Pairs: First, we zip nums2 and nums1 together to create pairs (nums2[i], nums1[i]), then sort these pairs in descending order by nums2 values. This ensures as we traverse from left to right, each nums2 value we encounter could serve as the minimum for our selection.

   nums = sorted(zip(nums2, nums1), reverse=True)

2. Initialize Data Structures:

   • q: A min heap to maintain the largest k elements from nums1
   • s: Running sum of elements currently in the heap
   • ans: Maximum score found so far

   q = []
   ans = s = 0

3. Traverse Sorted Pairs: For each pair (a, b) where a is from nums2 and b is from nums1:

   • Add b to the running sum and push it to the heap

   s += b
   heappush(q, b)

4. Maintain Heap Size and Calculate Score:

   • When the heap reaches size k, we have exactly k elements:
     • Calculate the score: s * a (sum of k largest nums1 values × current minimum a)
     • Update the maximum score
     • Remove the smallest element from the heap to make room for potentially better elements in future iterations

   if len(q) == k:
       ans = max(ans, s * a)
       s -= heappop(q)

The min heap ensures we always keep the k largest elements from nums1 seen so far. By processing in descending order of nums2, we guarantee that when we calculate the score at position i, the value nums2[i] is indeed the minimum among all selected indices, and we're pairing it with the optimal sum from nums1.

Time complexity: O(n log n) for sorting plus O(n log k) for heap operations. Space complexity: O(n) for the sorted array and O(k) for the heap.

Example Walkthrough

Let's walk through a concrete example with nums1 = [2, 1, 14, 12], nums2 = [11, 7, 13, 6], and k = 3.

Step 1: Create and sort pairs by nums2 (descending)

• Original pairs: (11,2), (7,1), (13,14), (6,12)
• Sorted pairs: (13,14), (11,2), (7,1), (6,12)

Step 2: Traverse sorted pairs with min heap

Iteration 1: Process (13,14)

• Add 14 to heap: heap = [14], sum = 14
• Heap size (1) < k (3), so no score calculation yet

Iteration 2: Process (11,2)

• Add 2 to heap: heap = [2, 14], sum = 16
• Heap size (2) < k (3), so no score calculation yet

Iteration 3: Process (7,1)

• Add 1 to heap: heap = [1, 2, 14], sum = 17
• Heap size equals k! Calculate score: 17 × 7 = 119
• Update ans = 119
• Remove smallest (1) from heap: heap = [2, 14], sum = 16

Iteration 4: Process (6,12)

• Add 12 to heap: heap = [2, 14, 12], sum = 28
• Heap size equals k! Calculate score: 28 × 6 = 168
• Update ans = 168
• Remove smallest (2) from heap: heap = [12, 14], sum = 26

Result: Maximum score is 168, achieved by selecting indices with values (13,14), (6,12), and (11,2) from the original arrays. The sum from nums1 is 14+12+2=28, the minimum from nums2 is min(13,6,11)=6, giving us 28×6=168.

The key insight: when we calculate the score at each position, that position's nums2 value is guaranteed to be the minimum (since we're traversing in descending order), and our heap maintains the k largest nums1 values we've seen so far that can be paired with this minimum.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n × log n)

The time complexity is dominated by two main operations:

1. Sorting the zipped pairs (nums2[i], nums1[i]) in descending order takes O(n × log n) time, where n is the length of the input arrays.
2. The for loop iterates through all n elements, and within each iteration:
   • heappush(q, b) takes O(log k) time (heap size is at most k)
   • heappop(q) takes O(log k) time when the heap size reaches k

Since k ≤ n, the loop operations contribute O(n × log k) which is bounded by O(n × log n). Therefore, the overall time complexity is O(n × log n).

Space Complexity: O(n)

The space complexity consists of:

1. The nums list storing the zipped and sorted pairs requires O(n) space.
2. The min-heap q stores at most k elements, requiring O(k) space.
3. Other variables (ans, s, a, b) use O(1) space.

Since k ≤ n, the total space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Heap Type Selection

Pitfall: Using a max heap instead of a min heap, or misunderstanding why we need a min heap.

Many developers intuitively think "we want maximum values, so let's use a max heap." However, we actually need a min heap to efficiently remove the smallest element when our heap exceeds size k.

Why this happens: When maintaining the k largest elements, we need to remove the smallest among them when adding a new element. A min heap gives us O(1) access to the smallest element.

Solution:

# Correct: Use min heap to maintain k largest elements
min_heap = []
heapq.heappush(min_heap, value)  # Smallest element at root

# Incorrect: Trying to use max heap (Python doesn't have built-in max heap)
# This would require negating values, adding complexity

2. Forgetting to Update Sum When Popping from Heap

Pitfall: Only popping from the heap without adjusting the running sum accordingly.

# Incorrect - sum becomes inaccurate
if len(min_heap) == k:
    max_score = max(max_score, current_sum * nums2_value)
    heapq.heappop(min_heap)  # Forgot to update current_sum!

# Correct - maintain accurate sum
if len(min_heap) == k:
    max_score = max(max_score, current_sum * nums2_value)
    current_sum -= heapq.heappop(min_heap)  # Update sum when removing

3. Calculating Score Before Having k Elements

Pitfall: Computing the score when heap size is less than k, which violates the problem constraint of selecting exactly k indices.

# Incorrect - calculates score even with fewer than k elements
for nums2_value, nums1_value in paired_values:
    current_sum += nums1_value
    heapq.heappush(min_heap, nums1_value)
    max_score = max(max_score, current_sum * nums2_value)  # Wrong!

# Correct - only calculate when we have exactly k elements
for nums2_value, nums1_value in paired_values:
    current_sum += nums1_value
    heapq.heappush(min_heap, nums1_value)
    if len(min_heap) == k:  # Check size first
        max_score = max(max_score, current_sum * nums2_value)
        current_sum -= heapq.heappop(min_heap)

4. Incorrect Pairing Order

Pitfall: Pairing the arrays incorrectly or sorting by the wrong array.

# Incorrect - sorting by nums1 instead of nums2
paired_values = sorted(zip(nums1, nums2), reverse=True)

# Correct - sort by nums2 (which will be our minimum)
paired_values = sorted(zip(nums2, nums1), reverse=True)

5. Edge Case: When n < k

Pitfall: Not considering that the arrays might have fewer elements than k.

Solution: The algorithm naturally handles this case because the heap never reaches size k, so no score is calculated and the default value of 0 is returned. However, it's worth being aware of this edge case for debugging purposes.