Medium

Choose K Elements With Maximum Sum

LAB

For each index, choose at most k nums2 values from prior indices whose nums1 value is smaller, and return the maximum possible sum for each index.

EXAMPLES

Example 1

Input
{
  "nums1": [
    4,
    2,
    1,
    5,
    3
  ],
  "nums2": [
    10,
    20,
    30,
    40,
    50
  ],
  "k": 2
}

Output
[
  80,
  30,
  0,
  80,
  50
]

FUNCTION SHAPE

nums1: intArraynums2: intArrayk: int→intArray

SOLUTION NOTE

The high level idea is if you consider nums1 in sorted order with index j, you know all previous indices j that we need to index into nums2 with. So basically you can use a heap of size k to maintain the sum of the max k. This is a min heap, so when we pop, we pop the smallest value out. This code is very common to maintaining a heap of size k:

pythonEXAMPLE

heappush(heap, nums2[sl[i][1]])
sum_heap += nums2[sl[i][1]]
if len(heap) == k+1:
    sum_heap -= heappop(heap)

Now given a certain index i, (in nums1 sorted order) we have the greedy optimal max sum of nums2 elements with these corresponding indices j.

What remains is re-assigning back to our original indices i. Sum_max_k indices are relative to the sorted order of nums1.

We need to binary search in the sorted list for (nums1[i], -inf). (see sorted list chapter) -1 to the index, and that gives us the sorted order index which we can index into for sum_max_k.

Reveal reference solution +

pythonREFERENCE

def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
    n = len(nums1)
    sl = sorted([(nums1[j], j) for j in range(n)])
    heap = []
    sum_max_k = [0] * n
    sum_heap = 0

    for i in range(n):
        heappush(heap, nums2[sl[i][1]])
        sum_heap += nums2[sl[i][1]]
        if len(heap) == k+1:
            sum_heap -= heappop(heap)
        sum_max_k[i] = sum_heap

    res = [0] * n

    for i in range(n):
        ii = bisect_left(sl, (nums1[i], -inf)) - 1
        if ii >= 0:
            res[i] = sum_max_k[ii]
    return res

TimeO(n log n)

SpaceO(n)

Open on LeetCode ↗