problemmediumalgorithmsleetcode-3478leetcode 3478leetcode3478

Choose K Elements With Maximum Sum

MediumUpdated: Jul 5, 2025
Practice on:
LeetCode

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.

Examples

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 <= 10^5
  • 1 <= nums1[i], nums2[i] <= 10^6
  • 1 <= k <= n

Solution

Method 1 – Sorting and Heap (Priority Queue)

Intuition

For each index i, we want to find all indices j where nums1[j] < nums1[i] and select the k largest nums2[j] values. This is a classic use case for sorting and a max-heap (priority queue).

Reasoning

By sorting the indices by nums1 value, we can process all indices in increasing order. For each i, all previous indices have nums1[j] < nums1[i]. We can maintain a max-heap of nums2[j] values for all such j and efficiently get the sum of the k largest values for each i.

Approach

  1. Create a list of indices sorted by nums1 value.
  2. For each index in sorted order, maintain a max-heap of nums2[j] for all previous indices.
  3. For each i, the answer is the sum of the k largest values in the heap (for all j with nums1[j] < nums1[i]).
  4. If there are fewer than k such values, sum as many as possible.

Edge cases:

  • If no j exists for i, answer is 0.
  • If fewer than k values, sum as many as possible.

Code

C++
class Solution {
public:
    vector<int> maxSum(vector<int>& nums1, vector<int>& nums2, int k) {
        int n = nums1.size();
        vector<int> idx(n);
        iota(idx.begin(), idx.end(), 0);
        sort(idx.begin(), idx.end(), [&](int a, int b) { return nums1[a] < nums1[b]; });
        vector<int> ans(n);
        multiset<int> heap;
        for (int i = 0; i < n; ++i) {
            int cur = idx[i];
            vector<int> topk;
            auto it = heap.rbegin();
            for (int cnt = 0; cnt < k && it != heap.rend(); ++cnt, ++it) topk.push_back(*it);
            ans[cur] = accumulate(topk.begin(), topk.end(), 0);
            heap.insert(nums2[cur]);
        }
        return ans;
    }
};
Java
class Solution {
    public int[] maxSum(int[] nums1, int[] nums2, int k) {
        int n = nums1.length;
        Integer[] idx = new Integer[n];
        for (int i = 0; i < n; i++) idx[i] = i;
        Arrays.sort(idx, Comparator.comparingInt(i -> nums1[i]));
        int[] ans = new int[n];
        PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder());
        for (int i = 0; i < n; i++) {
            int cur = idx[i];
            List<Integer> topk = new ArrayList<>();
            Iterator<Integer> it = heap.iterator();
            List<Integer> all = new ArrayList<>(heap);
            Collections.sort(all, Collections.reverseOrder());
            for (int cnt = 0; cnt < k && cnt < all.size(); cnt++) topk.add(all.get(cnt));
            ans[cur] = topk.stream().mapToInt(Integer::intValue).sum();
            heap.add(nums2[cur]);
        }
        return ans;
    }
}
Python
class Solution:
    def max_sum(self, nums1: list[int], nums2: list[int], k: int) -> list[int]:
        n = len(nums1)
        idx = sorted(range(n), key=lambda i: nums1[i])
        ans = [0] * n
        heap = []
        for i in range(n):
            cur = idx[i]
            topk = sorted(heap, reverse=True)[:k]
            ans[cur] = sum(topk)
            heap.append(nums2[cur])
        return ans
Go
func maxSum(nums1, nums2 []int, k int) []int {
    n := len(nums1)
    idx := make([]int, n)
    for i := range idx { idx[i] = i }
    sort.Slice(idx, func(i, j int) bool { return nums1[idx[i]] < nums1[idx[j]] })
    ans := make([]int, n)
    heap := []int{}
    for i := 0; i < n; i++ {
        cur := idx[i]
        tmp := append([]int{}, heap...)
        sort.Sort(sort.Reverse(sort.IntSlice(tmp)))
        sum := 0
        for cnt := 0; cnt < k && cnt < len(tmp); cnt++ {
            sum += tmp[cnt]
        }
        ans[cur] = sum
        heap = append(heap, nums2[cur])
    }
    return ans
}
Kotlin
class Solution {
    fun maxSum(nums1: IntArray, nums2: IntArray, k: Int): IntArray {
        val n = nums1.size
        val idx = nums1.indices.sortedBy { nums1[it] }
        val ans = IntArray(n)
        val heap = mutableListOf<Int>()
        for (i in 0 until n) {
            val cur = idx[i]
            val topk = heap.sortedDescending().take(k)
            ans[cur] = topk.sum()
            heap.add(nums2[cur])
        }
        return ans
    }
}
Rust
impl Solution {
    pub fn max_sum(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<i32> {
        let n = nums1.len();
        let mut idx: Vec<usize> = (0..n).collect();
        idx.sort_by_key(|&i| nums1[i]);
        let mut ans = vec![0; n];
        let mut heap = vec![];
        for &cur in &idx {
            let mut topk = heap.clone();
            topk.sort_unstable_by(|a, b| b.cmp(a));
            ans[cur] = topk.iter().take(k as usize).sum();
            heap.push(nums2[cur]);
        }
        ans
    }
}

Complexity

  • ⏰ Time complexity: O(n k log k), as for each index, we sort up to n elements and take the top k.
  • 🧺 Space complexity: O(n), as we store the heap and answer array.

Comments