3478. Choose K Elements With Maximum Sum LeetCode Solution

In this guide, you will get 3478. Choose K Elements With Maximum Sum LeetCode Solution with the best time and space complexity. The solution to Choose K Elements With Maximum Sum problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Problem Statement
Complexity Analysis
Choose K Elements With Maximum Sum solution in C++
Choose K Elements With Maximum Sum solution in Java
Choose K Elements With Maximum Sum solution in Python
Additional Resources

3478. Choose K Elements With Maximum Sum LeetCode Solution image

Problem Statement of Choose K Elements With Maximum Sum

You are given two integer arrays, nums1 and nums2, both of length n, along with a positive integer k.
For each index i from 0 to n – 1, perform the following:

Find all indices j where nums1[j] is less than nums1[i].
Choose at most k values of nums2[j] at these indices to maximize the total sum.

Return an array answer of size n, where answer[i] represents the result for the corresponding index i.

Example 1:

Input: nums1 = [4,2,1,5,3], nums2 = [10,20,30,40,50], k = 2
Output: [80,30,0,80,50]
Explanation:

For i = 0: Select the 2 largest values from nums2 at indices [1, 2, 4] where nums1[j] < nums1[0], resulting in 50 + 30 = 80.
For i = 1: Select the 2 largest values from nums2 at index [2] where nums1[j] < nums1[1], resulting in 30.
For i = 2: No indices satisfy nums1[j] < nums1[2], resulting in 0.
For i = 3: Select the 2 largest values from nums2 at indices [0, 1, 2, 4] where nums1[j] < nums1[3], resulting in 50 + 30 = 80.
For i = 4: Select the 2 largest values from nums2 at indices [1, 2] where nums1[j] < nums1[4], resulting in 30 + 20 = 50.

Example 2:

Input: nums1 = [2,2,2,2], nums2 = [3,1,2,3], k = 1
Output: [0,0,0,0]
Explanation:
Since all elements in nums1 are equal, no indices satisfy the condition nums1[j] < nums1[i] for any i, resulting in 0 for all positions.

Constraints:

n == nums1.length == nums2.length
1 <= n <= 105
1 <= nums1[i], nums2[i] <= 106
1 <= k <= n

Complexity Analysis

Time Complexity: O(n\log k)
Space Complexity: O(n + k)

3478. Choose K Elements With Maximum Sum LeetCode Solution in C++

class Solution {
 public:
  vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
    const int n = nums1.size();
    vector<long long> ans(n);
    vector<pair<int, int>> numAndIndexes;
    priority_queue<long long, vector<long long>, greater<long long>> minHeap;

    for (int i = 0; i < n; i++)
      numAndIndexes.emplace_back(nums1[i], i);

    ranges::sort(numAndIndexes);

    const int firstIndex = numAndIndexes[0].second;
    minHeap.push(nums2[firstIndex]);
    long sum = nums2[firstIndex];

    for (int i = 1; i < n; ++i) {
      const auto& [currNum, currIndex] = numAndIndexes[i];
      const auto& [prevNum, prevIndex] = numAndIndexes[i - 1];
      if (currNum == prevNum)
        ans[currIndex] = ans[prevIndex];
      else
        ans[currIndex] = sum;
      minHeap.push(nums2[currIndex]);
      sum += nums2[currIndex];
      if (minHeap.size() == k + 1)
        sum -= minHeap.top(), minHeap.pop();
    }

    return ans;
  }
};
/* code provided by PROGIEZ */

3478. Choose K Elements With Maximum Sum LeetCode Solution in Java

class Solution {
  public long[] findMaxSum(int[] nums1, int[] nums2, int k) {
    final int n = nums1.length;
    long[] ans = new long[n];
    Pair<Integer, Integer>[] numAndIndexes = new Pair[n];
    Queue<Long> minHeap = new PriorityQueue<>();

    for (int i = 0; i < n; ++i)
      numAndIndexes[i] = new Pair<>(nums1[i], i);

    Arrays.sort(numAndIndexes, Comparator.comparing(Pair::getKey));

    final int firstIndex = numAndIndexes[0].getValue();
    minHeap.offer((long) nums2[firstIndex]);
    long sum = nums2[firstIndex];

    for (int i = 1; i < n; ++i) {
      final int currNum = numAndIndexes[i].getKey();
      final int currIndex = numAndIndexes[i].getValue();
      final int prevNum = numAndIndexes[i - 1].getKey();
      final int prevIndex = numAndIndexes[i - 1].getValue();
      if (currNum == prevNum)
        ans[currIndex] = ans[prevIndex];
      else
        ans[currIndex] = sum;
      minHeap.offer((long) nums2[currIndex]);
      sum += nums2[currIndex];
      if (minHeap.size() == k + 1)
        sum -= minHeap.poll();
    }

    return ans;
  }
}
// code provided by PROGIEZ

3478. Choose K Elements With Maximum Sum LeetCode Solution in Python

class Solution:
  def findMaxSum(self, nums1: list[int], nums2: list[int], k: int) -> list[int]:
    ans = [0] * len(nums1)
    numAndIndexes = sorted([(num, i) for i, num in enumerate(nums1)])
    minHeap = []

    firstIndex = numAndIndexes[0][1]
    heapq.heappush(minHeap, nums2[firstIndex])
    summ = nums2[firstIndex]

    for (prevNum, prevIndex), (currNum, currIndex) in itertools.pairwise(numAndIndexes):
      if currNum == prevNum:
        ans[currIndex] = ans[prevIndex]
      else:
        ans[currIndex] = summ
      heapq.heappush(minHeap, nums2[currIndex])
      summ += nums2[currIndex]
      if len(minHeap) == k + 1:
        summ -= heapq.heappop(minHeap)

    return ans
 # code by PROGIEZ