Skip to main content
Talentd

Choose K Elements With Maximum Sum - Solution & Explanation

MediumArraySortingHeap (Priority Queue)10 min readAsked at: Bloomberg
Practice this problem

Problem Statement

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

Approach Overview

Problem Overview: You are given two arrays where each index represents a pair of values. For every position i, compute the maximum possible sum of up to k values from other indices whose first value is strictly smaller than nums1[i]. The goal is to efficiently track the best k candidates while scanning the array.

Approach 1: Brute Force with Heap per Index (O(n² log k) time, O(k) space)

For each index i, iterate through the entire array and collect elements where nums1[j] < nums1[i]. Maintain a min-heap of size at most k containing the largest nums2[j] values seen so far. Each candidate pushes into the heap and the smallest element is removed if the size exceeds k. The sum of the heap gives the best result for that index. This method is straightforward but expensive because the full scan repeats for every element.

Approach 2: Sorting + Priority Queue (Min-Heap) (O(n log n + n log k) time, O(n + k) space)

Sort indices by nums1 so elements with smaller values are processed first. As you iterate through the sorted order, maintain a running min-heap storing the largest k values from nums2 that have already been seen. Also track the current sum of the heap. For each group of indices with the same nums1 value, compute their answers using the current sum before inserting their nums2 values into the heap. This ensures only elements with strictly smaller nums1 contribute to the result. Each insertion or removal from the heap costs O(log k), keeping the overall runtime efficient.

The key insight: sorting transforms the "smaller value" condition into a prefix problem. Once elements are processed in increasing order, every previously visited element automatically satisfies the constraint.

This pattern appears frequently when combining sorting with incremental state tracking using a heap (priority queue). The input itself is just a standard array, but ordering it unlocks a much more efficient solution.

Recommended for interviews: The sorting + min-heap approach is the expected solution. It demonstrates recognition of ordering constraints and efficient top-k maintenance. Mentioning the brute force approach first shows you understand the problem space, but implementing the O(n log n + n log k) method proves strong algorithmic judgment.

Solution

We can convert the array nums1 into an array arr, where each element is a tuple (x, i), representing the value x at index i in nums1. Then, we sort the array arr in ascending order by x.

We use a min-heap pq to maintain the elements from the array nums2. Initially, pq is empty. We use a variable s to record the sum of the elements in pq. Additionally, we use a pointer j to maintain the current position in the array arr that needs to be added to pq.

We traverse the array arr. For the h-th element (x, i), we add all elements nums2[arr[j][1]] to pq that satisfy j < h and arr[j][0] < x, and add these elements to s. If the size of pq exceeds k, we pop the smallest element from pq and subtract it from s. Then, we update the value of ans[i] to s.

After traversing, we return the answer array ans.

The time complexity is O(n log n), and the space complexity is O(n). Here, n is the length of the array.

Code

PythonJavaC++GoTypeScript

Python

Java

C++

Go

TypeScript

Try this approach in the editor →

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
Brute Force with Heap per IndexO(n² log k)O(k)Useful for understanding the problem or when n is very small
Sorting + Min-Heap (Priority Queue)O(n log n + n log k)O(n + k)Best general solution; efficiently maintains top k candidates while processing sorted elements

Video Solution

Choose K Elements With Maximum Sum | Brute Force | Optimal Leetcode 3478 | codestorywithMIK • codestorywithMIK • 7,067 views views

Watch 7 more video solutions →

Frequently Asked Questions

Is Choose K Elements With Maximum Sum easy or hard?
Choose K Elements With Maximum Sum is generally considered a medium difficulty problem. The challenge comes from recognizing that sorting converts the constraint into a prefix problem and that a heap can efficiently maintain the best k candidates.
Choose K Elements With Maximum Sum Python/Java solution
Most implementations follow the same structure: sort indices by nums1, iterate through them, maintain a min-heap of size k for nums2 values, and track a running sum. Python typically uses heapq, while Java implementations use PriorityQueue with O(log k) insertion and removal.
How to solve Choose K Elements With Maximum Sum efficiently?
Sort indices by nums1 so that all valid contributors appear before the current element. Maintain a running min-heap containing the largest k values from nums2 seen so far and track their sum. For each index, use the current heap sum as the answer before inserting its nums2 value into the heap.
What is the best approach for Choose K Elements With Maximum Sum?
The optimal solution sorts indices by the first array and processes them in increasing order while maintaining a min-heap of the largest k values from the second array. This allows efficient tracking of the best candidates that satisfy the "smaller value" condition. The time complexity is O(n log n + n log k) and space complexity is O(n + k).
Is Choose K Elements With Maximum Sum asked at Google/Amazon/Meta?
Problems combining sorting with heap-based top-k tracking frequently appear in interviews at companies like Amazon, Google, and Meta. Variants involving maintaining the k largest elements under ordering constraints are especially common in system design and algorithm screening rounds.
What data structure is used in Choose K Elements With Maximum Sum?
The core data structure is a min-heap (priority queue). It keeps the largest k values from nums2 while allowing efficient removal of the smallest when the heap exceeds size k. Sorting is also required to enforce the nums1 ordering constraint.
What is the time complexity of Choose K Elements With Maximum Sum?
The optimal algorithm runs in O(n log n + n log k). Sorting the indices by nums1 takes O(n log n), and each heap insertion or removal while maintaining the top k elements costs O(log k). Since each element is processed once, the total heap work is O(n log k).

Ready to solve this problem?

Practice Choose K Elements With Maximum Sum with our built-in code editor and test cases.

Practice on FleetCode
Previous

Two Sum

Next

Median of Two Sorted Arrays