Choose K Elements With Maximum Sum - Practice Coding Problems

Choose K Elements With Maximum Sum - Problem

Choose K Elements With Maximum Sum

You're given two integer arrays nums1 and nums2, both of length n, along with a positive integer k. Your task is to process each index and find the maximum possible sum by strategically selecting elements.

For each index i from 0 to n-1, you need to:
1. Find all indices j where nums1[j] < nums1[i]
2. From these valid indices, choose at most k values from nums2[j] to maximize the total sum

Return an array answer of size n, where answer[i] represents the maximum sum achievable for index i.

Example:
If nums1 = [3, 1, 4, 2], nums2 = [10, 20, 5, 15], and k = 2:
- For index 0 (nums1[0] = 3): indices with smaller values are [1] → choose nums2[1] = 20
- For index 2 (nums1[2] = 4): indices with smaller values are [0,1,3] → choose top 2: nums2[1] + nums2[3] = 35

Input & Output

example_1.py — Basic Case

$ Input: nums1 = [3, 1, 4, 2], nums2 = [10, 20, 5, 15], k = 2

› Output: [20, 0, 35, 20]

💡 Note: For i=0 (nums1[0]=3): valid indices are [1] where nums1[1]=1<3, so max sum = nums2[1]=20. For i=1 (nums1[1]=1): no valid indices since no element < 1, so sum = 0. For i=2 (nums1[2]=4): valid indices are [0,1,3], top 2 values from nums2 are 20+15=35. For i=3 (nums1[3]=2): valid index is [1], so sum = nums2[1]=20.

example_2.py — Small k

$ Input: nums1 = [5, 2, 8, 1], nums2 = [1, 3, 2, 4], k = 1

› Output: [3, 4, 4, 0]

💡 Note: For i=0: valid indices [1,3], max value is max(3,4)=4. Wait, let me recalculate: For i=0 (nums1[0]=5): valid indices where nums1[j]<5 are [1,2,3] but nums1[2]=8>5, so valid are [1,3], max from nums2[1],nums2[3] is max(3,4)=4. Actually for i=0: nums1[1]=2<5✓, nums1[3]=1<5✓, nums1[2]=8≥5✗, so nums2[1]=3, nums2[3]=4, max is 4.

example_3.py — Edge Case All Same

$ Input: nums1 = [5, 5, 5], nums2 = [1, 2, 3], k = 2

› Output: [0, 0, 0]

💡 Note: Since all elements in nums1 are equal (5), for any index i, there are no indices j where nums1[j] < nums1[i]. Therefore, no valid elements can be chosen, and the sum for each index is 0.

Constraints

1 ≤ n ≤ 1000
1 ≤ k ≤ n
-10⁴ ≤ nums1[i], nums2[i] ≤ 10⁴
nums1 and nums2 have the same length n

Visualization

Tap to expand

$35⚡ Efficient Scouting with Priority QueueInstead of sorting all salaries every time:• Use a min-heap of size k to track top salaries• Add salary if heap not full, or replace minimum if new salary is higher• This reduces time complexity from O(n² log n) to O(n² log k)• Perfect for when k is small compared to n!💡 Key Insight: Use heap to avoid repeatedly sorting the same elements

Understanding the Visualization

Identify the Star Player

Select the player you're considering (index i) and note their skill level

Find Eligible Recruits

Look at all other players and identify those with lower skill ratings

Optimize Budget Allocation

From eligible players, select the top-k highest salaries to maximize investment

Record Total Investment

Sum the selected salaries and record as the result for this star player

Key Takeaway

🎯 Key Insight: The problem is essentially finding top-k elements from a filtered subset for each position. Using a min-heap of size k avoids the overhead of full sorting and provides optimal performance, especially when k is much smaller than n.

Asked in

G Google 42 a Amazon 38 f Meta 29 ⊞ Microsoft 24

The optimal approach uses a min-heap of size k for each index to efficiently maintain the top-k largest elements from valid indices. For each index i, we identify all indices j where nums1[j] < nums1[i], then use a heap to track the k largest values from the corresponding nums2[j] positions. This avoids the overhead of full sorting and achieves O(n² log k) time complexity, which is better than the brute force O(n² log n) approach, especially when k is small.

Common Approaches

Approach	Time	Space	Notes
✓ Brute Force (Nested Loops)	O(n² log n)	O(n)	For each index, find all valid indices and sort to get top-k elements
Sorting + Priority Queue (Optimal)	O(n² log k)	O(n + k)	Use sorted indices and max-heap to efficiently find top-k elements

Brute Force (Nested Loops) — Algorithm Steps

For each index i from 0 to n-1
Initialize an empty list to store valid nums2 values
Iterate through all indices j to find nums1[j] < nums1[i]
Collect nums2[j] for valid indices j
Sort collected values in descending order
Sum the first min(k, length) elements
Store result in answer[i]

Visualization

Tap to expand

Step-by-Step Walkthrough

Select Current Index

Choose index i and identify nums1[i] as the threshold

Find Valid Indices

Scan all indices j where nums1[j] < nums1[i]

Collect and Sort

Gather nums2[j] values and sort them in descending order

Sum Top K

Take the first k elements and sum them

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>

int compare(const void *a, const void *b) {
    return (*(int*)b - *(int*)a); // Descending order
}

int* solution(int* nums1, int* nums2, int n, int k, int* returnSize) {
    int* answer = (int*)malloc(n * sizeof(int));
    *returnSize = n;
    
    for (int i = 0; i < n; i++) {
        int* validValues = (int*)malloc(n * sizeof(int));
        int count = 0;
        
        for (int j = 0; j < n; j++) {
            if (nums1[j] < nums1[i]) {
                validValues[count++] = nums2[j];
            }
        }
        
        // Sort in descending order
        qsort(validValues, count, sizeof(int), compare);
        
        int maxSum = 0;
        int limit = (k < count) ? k : count;
        for (int idx = 0; idx < limit; idx++) {
            maxSum += validValues[idx];
        }
        answer[i] = maxSum;
        
        free(validValues);
    }
    
    return answer;
}

int main() {
    // Parse input and call solution
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n² log n)

For each of n elements, we potentially check all n elements (O(n²)), and sort up to n elements (O(n log n))

⚠ Quadratic Growth

Space Complexity

O(n)

Additional space for storing valid elements and the result array

⚡ Linearithmic Space

34.2K Views

Medium-High Frequency

~18 min Avg. Time

1.2K Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Brute Force (Nested Loops) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler