Choose K Elements With Maximum Sum

Choose K Elements With Maximum Sum - Problem

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.

Input & Output

Example 1 — Basic Case

$ Input: nums1 = [2,4,1,0,3], nums2 = [3,4,6,2,1], k = 3

› Output: [8,11,2,0,11]

💡 Note: For index 0 (nums1[0]=2): valid indices are 2,3 where nums1 values are 1,0 (both < 2). Corresponding nums2 values are [6,2]. Sum = 8. For index 1 (nums1[1]=4): valid indices are 0,2,3,4 with nums2 values [3,6,2,1]. Top 3 are [6,3,2] with sum 11.

Example 2 — Small k

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

› Output: [0,0,9]

💡 Note: For index 0: no elements < 5, sum = 0. For index 1: no elements < 1, sum = 0. For index 2: only index 1 has nums1[1]=1 < 3, so sum = nums2[1] = 9.

Example 3 — All Elements Valid

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

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

💡 Note: For index 3 (nums1[3]=4): valid indices are 0,1,2 with nums1 values 1,2,3 (all < 4). Corresponding nums2 values are [4,3,2]. Top 2 are [4,3] with sum 7.

Constraints

1 ≤ n ≤ 1000
1 ≤ k ≤ n
-1000 ≤ nums1[i], nums2[i] ≤ 1000

Visualization

Tap to expand

Asked in

G Google 23 a Amazon 18 f Facebook 15 A Apple 12

The key insight is to efficiently find and select the top k elements from valid candidates for each index. The brute force approach works by sorting candidates for each query, taking O(n² log n). The heap approach optimizes this to O(n² log k) by maintaining only k elements. Best approach depends on k vs n ratio. Time: O(n² log k), Space: O(k).

Common Approaches

✓ Brute Force with Sorting

⏱️ Time: O(n² log n) Space: O(n)

For each index i, iterate through all indices j to find where nums1[j] < nums1[i]. Collect corresponding nums2[j] values, sort them in descending order, and sum the top k elements.

Priority Queue (Heap) Optimization

⏱️ Time: O(n² log k) Space: O(k)

For each index, use a priority queue to maintain the k largest elements from valid indices. This avoids sorting the entire candidate list repeatedly.

Sort with Indices Optimization

⏱️ Time: O(n² log n) Space: O(n)

Create sorted pairs of (nums1[i], nums2[i], i) and use two pointers or binary search to efficiently find valid elements for each query. This reduces the time complexity for finding candidates.

Brute Force with Sorting — Algorithm Steps

For each index i, scan all indices j
Collect nums2[j] where nums1[j] < nums1[i]
Sort collected values in descending order
Sum the top k values

Visualization

Tap to expand

Step-by-Step Walkthrough

Find Candidates

For index i, find all j where nums1[j] < nums1[i]

Sort Values

Sort corresponding nums2[j] values in descending order

Sum Top K

Take sum of first k elements

Code -

solution.c — C

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

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

int parseArray(char* line, int* arr) {
    int count = 0;
    int i = 0;
    int len = strlen(line);
    
    // Skip opening bracket
    while (i < len && line[i] != '[') i++;
    i++;
    
    while (i < len && line[i] != ']') {
        if (line[i] == ',' || line[i] == ' ') {
            i++;
            continue;
        }
        
        // Parse number
        int num = 0;
        int negative = 0;
        if (line[i] == '-') {
            negative = 1;
            i++;
        }
        
        while (i < len && line[i] >= '0' && line[i] <= '9') {
            num = num * 10 + (line[i] - '0');
            i++;
        }
        
        if (negative) num = -num;
        arr[count++] = num;
    }
    
    return count;
}

int* solution(int* nums1, int* nums2, int n, int k, int* returnSize) {
    *returnSize = n;
    int* answer = (int*)calloc(n, sizeof(int));
    
    for (int i = 0; i < n; i++) {
        int* candidates = (int*)malloc(n * sizeof(int));
        int count = 0;
        
        for (int j = 0; j < n; j++) {
            if (nums1[j] < nums1[i]) {
                candidates[count++] = nums2[j];
            }
        }
        
        qsort(candidates, count, sizeof(int), compare);
        
        int total = 0;
        int limit = (k < count) ? k : count;
        for (int idx = 0; idx < limit; idx++) {
            total += candidates[idx];
        }
        answer[i] = total;
        
        free(candidates);
    }
    
    return answer;
}

int main() {
    char line[10000];
    int nums1[1000], nums2[1000];
    int n1, n2, k;
    
    // Read nums1
    fgets(line, sizeof(line), stdin);
    n1 = parseArray(line, nums1);
    
    // Read nums2
    fgets(line, sizeof(line), stdin);
    n2 = parseArray(line, nums2);
    
    // Read k
    scanf("%d", &k);
    
    int returnSize;
    int* result = solution(nums1, nums2, n1, k, &returnSize);
    
    printf("[");
    for (int i = 0; i < returnSize; i++) {
        printf("%d", result[i]);
        if (i < returnSize - 1) printf(",");
    }
    printf("]\n");
    
    free(result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n² log n)

For each of n indices, we potentially collect n elements and sort them (O(n log n))

⚠ Quadratic Growth

Space Complexity

O(n)

Temporary array to store candidates for each index can hold up to n elements

⚡ Linearithmic Space

34.1K Views

Medium Frequency

~25 min Avg. Time

892 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 with Sorting — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler