3478. Choose K Elements With Maximum Sum

Problem Description

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.

Intuition

To solve the problem, we need to identify indices j such that nums1[j] is less than nums1[i] for each i. Among these indices, we want to select at most k corresponding values from nums2 to maximize their total sum.

A direct approach would involve iterating over each index and comparing every other index, which could be inefficient for large arrays. Therefore, we employ an optimized approach using sorting and a priority queue:

Sorting Tuple Array: First, we convert nums1 into an array of tuples arr, where each tuple consists of a value from nums1 and its index. We then sort this array based on the values in nums1.
Maintaining Maximum Sum with Min-Heap: As we traverse the sorted array, we use a min-heap (priority queue) to keep track of the largest k sums for the values in nums2. For each element x in the sorted array, we consider previous elements that have a lower value (nums1[j] < nums1[i]) and add their corresponding nums2 values to the heap.
Updating the Answer: For each index i after processing it in the sorted order, we store the sum of the k largest elements from nums2 (selected using the heap) in the answer at index i.

By maintaining only the k largest sums efficiently with a min-heap, we can ensure that our solution is optimal and runs faster than a naive approach.

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

Solution Approach

Sorting + Priority Queue (Min-Heap)

We will utilize sorting and a priority queue (min-heap) to efficiently solve the problem.

Convert and Sort Array:
- Construct a new array of tuples arr, where each tuple contains a value from nums1 and its corresponding index: arr = [(x, i) for i, x in enumerate(nums1)].
- Sort the array arr based on the values in nums1.
Initialize Data Structures:
- Use a min-heap pq to keep track of values from nums2 corresponding to indices where values from nums1 are less than the current element.
- Use a variable s to keep track of the sum of elements in the heap.
- Initialize a pointer j to iterate over elements for each current element in nums1.
- Create an answer array ans of size n initialized to zero.
Process Each Element:
- Iterate over each element (x, i) in the sorted array arr.
- For every element where arr[j][0] < x, push the corresponding nums2 value into the heap pq.
- Add this value to the sum s.
- If the size of pq exceeds k, remove the smallest element from the heap using heappop, and subtract its value from s.
- Once all valid previous elements are considered, assign the current sum s to ans[i].
Complete and Return:
- After processing all elements, return the ans array, which contains the maximum sum for each i.

This approach efficiently computes the solution by leveraging sorting and a priority queue to manage and optimize the selection of elements in nums2, achieving the desired maximal sum for each index i.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's take a simple example to illustrate the solution approach.

Consider nums1 = [3, 1, 4], nums2 = [2, 5, 6], and k = 1.

We need to find the indices j where nums1[j] is less than nums1[i] for each i, and choose at most k values from nums2[j] to maximize the sum.

Convert and Sort Array:
- Construct the array of tuples: arr = [(3, 0), (1, 1), (4, 2)].
- Sort arr based on nums1 values: arr = [(1, 1), (3, 0), (4, 2)].
Initialize Data Structures:
- Use a min-heap pq and a variable s as an initial sum which is 0.
- Initialize a pointer j at 0 for iteration.
- Create an answer array ans of size 3 initialized to zero: ans = [0, 0, 0].
Process Each Element:
- For (1, 1): No previous elements in nums1 are smaller, so ans[1] = 0.
- For (3, 0):
  - arr[j][0] = 1 < 3 (valid j), add nums2[1] = 5 to heap: pq = [5], s = 5.
  - Since k = 1, no elements are removed from the heap.
  - Assign ans[0] = 5.
- For (4, 2):
  - arr[j][0] = 3 < 4, nums2[0] = 2 is added to heap: pq = [2, 5], s = 7.
  - The heap size exceeds k, remove the smallest element 2: pq = [5], s = 5.
  - Assign ans[2] = 5.
Complete and Return:
- The final ans array is [5, 0, 5].
- This means for each index i, the maximum sum of values from nums2 based on constraints is stored in ans[i].

The output [5, 0, 5] is achieved through sorting, and using a priority queue maintaining efficient management of values from nums2.

Solution Implementation

1from typing import List
2import heapq
3
4class Solution:
5    def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
6        # Create a list of tuples where each tuple is (value from nums1, index)
7        indexed_nums = [(value, index) for index, value in enumerate(nums1)]
8        # Sort the list based on the value from nums1
9        indexed_nums.sort()
10      
11        # Min-heap to store elements from nums2
12        min_heap = []
13        # To accumulate the sum of values from the heap
14        current_sum = 0
15        # Pointer for the sorted array
16        start_pointer = 0
17        # Length of nums1
18        n = len(indexed_nums)
19        # Resultant list to store the maximum sums
20        max_sums = [0] * n
21      
22        # Iterate through the sorted tuples
23        for end_pointer, (value, original_index) in enumerate(indexed_nums):
24            # Add elements to the heap and adjust start_pointer
25            while start_pointer < end_pointer and indexed_nums[start_pointer][0] < value:
26                corresponding_value = nums2[indexed_nums[start_pointer][1]]
27                heapq.heappush(min_heap, corresponding_value)
28                current_sum += corresponding_value
29                # Ensure the heap size does not exceed k
30                if len(min_heap) > k:
31                    current_sum -= heapq.heappop(min_heap)
32                start_pointer += 1
33          
34            # Record the current sum of the k largest elements
35            max_sums[original_index] = current_sum
36      
37        return max_sums
38

1import java.util.Arrays;
2import java.util.PriorityQueue;
3
4class Solution {
5    public long[] findMaxSum(int[] nums1, int[] nums2, int k) {
6        int n = nums1.length;
7        int[][] pairedArray = new int[n][2]; // Create a 2D array to store pairs of elements and their indices
8
9        // Populate the pairedArray with elements and their indices from nums1
10        for (int i = 0; i < n; ++i) {
11            pairedArray[i] = new int[] {nums1[i], i};
12        }
13
14        // Sort the pairs based on the first element of each pair
15        Arrays.sort(pairedArray, (a, b) -> a[0] - b[0]);
16
17        // Min-heap priority queue to keep track of the smallest elements
18        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
19        long currentSum = 0;
20        long[] result = new long[n];
21        int j = 0;
22
23        // Iterate over each element in the sorted array
24        for (int h = 0; h < n; ++h) {
25            int currentNum = pairedArray[h][0];
26            int indexInOriginalArray = pairedArray[h][1];
27
28            // Add elements from nums2 to the priority queue and update the sum until currentNum is reached
29            while (j < h && pairedArray[j][0] < currentNum) {
30                int numFromSecondArray = nums2[pairedArray[j][1]];
31                minHeap.offer(numFromSecondArray);
32                currentSum += numFromSecondArray;
33
34                // If the heap size exceeds k, remove the smallest element and update the sum
35                if (minHeap.size() > k) {
36                    currentSum -= minHeap.poll();
37                }
38                ++j;
39            }
40
41            // Store the current sum in the result array at the original index
42            result[indexInOriginalArray] = currentSum;
43        }
44
45        return result;
46    }
47}
48

1#include <vector>
2#include <queue>
3#include <algorithm>
4
5using namespace std;
6
7class Solution {
8public:
9    // Function to find the maximum sum for each element in nums1
10    vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
11        int n = nums1.size(); // Size of the input vectors
12        vector<pair<int, int>> indexedNums1(n);
13
14        // Create a vector of pairs where each pair contains an element from nums1 and its index
15        for (int i = 0; i < n; ++i) {
16            indexedNums1[i] = {nums1[i], i};
17        }
18
19        // Sort the pairs in ascending order based on the element from nums1
20        sort(indexedNums1.begin(), indexedNums1.end());
21
22        // Min-heap to maintain the smallest elements from nums2
23        priority_queue<int, vector<int>, greater<int>> minHeap;
24
25        long long currentSum = 0; // Current sum of elements taken from nums2
26        int j = 0;  // Pointer to track the elements considered from nums1
27        vector<long long> result(n); // Vector to store the results
28
29        // Iterate over each element in the sorted indexedNums1
30        for (int h = 0; h < n; ++h) {
31            // Destructure the current pair into value and original index
32            auto [x, i] = indexedNums1[h];
33
34            // Add elements from nums2 to the min-heap while they are smaller than the current element x from nums1
35            while (j < h && indexedNums1[j].first < x) {
36                int y = nums2[indexedNums1[j].second];
37                minHeap.push(y); // Add element from nums2 to the min-heap
38                currentSum += y; // Update the current sum
39
40                // Ensure the heap size does not exceed k and adjust the sum if necessary
41                if (minHeap.size() > k) {
42                    currentSum -= minHeap.top();
43                    minHeap.pop();
44                }
45                ++j;
46            }
47            // Store the current sum for the original index of the element
48            result[i] = currentSum;
49        }
50        return result;
51    }
52};
53

1// This function finds the maximum sum of elements from two arrays within certain conditions.
2function findMaxSum(nums1: number[], nums2: number[], k: number): number[] {
3    const n = nums1.length; // Length of nums1 (and nums2, assuming equal lengths)
4  
5    // Pair elements of nums1 with their indices and sort based on values in ascending order
6    const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]);
7  
8    // Initialize a min-priority queue to keep track of the k largest elements in nums2
9    const pq = new MinPriorityQueue();
10  
11    let s = 0;  // Sum of the largest elements tracked in the queue
12    let j = 0;  // Pointer for managing the queue
13
14    const ans: number[] = Array(k).fill(0); // Result array to store maximum sums
15
16    // Iterate over sorted elements of nums1
17    for (let h = 0; h < n; ++h) {
18        const [x, i] = arr[h]; // Current value from nums1 and its original index
19      
20        // Maintain the min-priority queue with the k largest corresponding nums2 values
21        while (j < h && arr[j][0] < x) {
22            const y = nums2[arr[j++][1]]; // Get corresponding value from nums2
23            pq.enqueue(y);  // Add it to the priority queue
24            s += y; // Increment the sum
25
26            // If the queue exceeds size k, remove the smallest element
27            if (pq.size() > k) {
28                s -= pq.dequeue();
29            }
30        }
31      
32        // Store the sum of the k (or fewer) largest elements in nums2 for the current index
33        ans[i] = s;
34    }
35
36    return ans; // Return the array of maximum sums
37}
38

Time and Space Complexity

The code sorts the arr list, which takes O(n \log n) time. Then, it iterates over the array and performs heap operations on the priority queue pq. The operations on the heap (insertion and extraction) take O(\log k) time, and they occur at most n times, where k is the maximum heap size. Thus, the overall time complexity of the code is O(n \log n).

The space complexity is O(n) since we are storing the arr array and a priority queue that can grow to size k, and k can potentially be O(n). Therefore, the space required is proportional to n.

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

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan: