Choose K Elements With Maximum Sum
MedPrompt
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
jwherenums1[j]is less thannums1[i]. - Choose at most
kvalues ofnums2[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 fromnums2at indices[1, 2, 4]wherenums1[j] < nums1[0], resulting in50 + 30 = 80. - For
i = 1: Select the 2 largest values fromnums2at index[2]wherenums1[j] < nums1[1], resulting in 30. - For
i = 2: No indices satisfynums1[j] < nums1[2], resulting in 0. - For
i = 3: Select the 2 largest values fromnums2at indices[0, 1, 2, 4]wherenums1[j] < nums1[3], resulting in50 + 30 = 80. - For
i = 4: Select the 2 largest values fromnums2at indices[1, 2]wherenums1[j] < nums1[4], resulting in30 + 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.length1 <= n <= 1051 <= nums1[i], nums2[i] <= 1061 <= k <= n
Approaches
3 approaches with complexity analysis and trade-offs.
This approach directly translates the problem description into a straightforward algorithm. For each index i, it scans the entire nums1 array to find all indices j that satisfy the condition nums1[j] < nums1[i]. The corresponding nums2[j] values are collected. To efficiently find the k largest values from this collection, a min-priority queue (min-heap) of size k is used. This heap keeps track of the k largest elements seen so far for the current i. Finally, the elements in the heap are summed up to get the result for answer[i].
Algorithm
- Initialize an empty
answerarray of sizen. - Iterate through each index
ifrom0ton-1.- For each
i, create a temporary list to store candidatenums2values. - Iterate through all indices
jfrom0ton-1. - If
nums1[j] < nums1[i], addnums2[j]to the list of candidates. - To find the sum of the top
kcandidates, use a min-priority queue of sizek. - Iterate through the candidates:
- Add the candidate value to the priority queue.
- If the priority queue size exceeds
k, remove the smallest element (the root).
- After checking all candidates, sum the elements remaining in the priority queue.
- Store this sum in
answer[i].
- For each
- Return the
answerarray.
Walkthrough
The brute-force method involves a nested loop structure. The outer loop iterates through each element of nums1 to compute the corresponding answer[i]. The inner loop iterates through all other elements to find potential candidates based on the condition nums1[j] < nums1[i]. For each i, we gather all valid nums2[j] values. Instead of sorting this list of candidates every time, which would be O(N log N), we can optimize by using a min-heap of size k. As we find candidates, we add them to the heap. If the heap's size grows beyond k, we remove the smallest element. This ensures the heap always contains the top k candidates encountered so far. The sum of these k elements gives us the desired result for index i.
import java.util.PriorityQueue; class Solution { public long[] getResults(int[] nums1, int[] nums2, int k) { int n = nums1.length; long[] answer = new long[n]; for (int i = 0; i < n; i++) { PriorityQueue<Integer> minHeap = new PriorityQueue<>(); for (int j = 0; j < n; j++) { if (nums1[j] < nums1[i]) { minHeap.offer(nums2[j]); if (minHeap.size() > k) { minHeap.poll(); } } } long currentSum = 0; for (int val : minHeap) { currentSum += val; } answer[i] = currentSum; } return answer; }}Complexity
Time
O(n^2 * log k) The outer loop runs `n` times for each index `i`. The inner loop also runs `n` times to find candidates. Inside the inner loop, each heap operation (offer/poll) takes `O(log k)` time. This results in a total time complexity of `O(n * n * log k)`.
Space
O(k) For each `i`, we use a priority queue that stores at most `k` elements. The `answer` array requires `O(n)` space, which is part of the output.
Trade-offs
Pros
Simple to understand and implement.
Directly follows the problem statement.
Cons
Extremely inefficient due to nested loops.
Will result in a 'Time Limit Exceeded' error on platforms for the given constraints.
This approach significantly improves performance by sorting the elements and using a powerful data structure, the Fenwick Tree (or Binary Indexed Tree). We first sort the elements based on nums1. This allows us to process queries in an order where the set of candidate nums2 values grows incrementally. To handle the nums2 values, which can be large, we use coordinate compression to map them to a smaller range of ranks. Two Fenwick Trees are used: one to maintain the count of numbers added and another for their sum. For each query, we find the k-th largest element by performing a binary search over the ranks, using the count BIT to guide the search. This allows us to find the sum of the top k elements much faster than the brute-force method.
Algorithm
- Coordinate Compression: Create a sorted list of unique values from
nums2and a map to get the rank of each value. Letmbe the number of unique values. - Combine and Sort: Create tuples of
(nums1[i], nums2[i], original_index_i)and sort these tuples based onnums1in ascending order. - Initialize: Set up two Fenwick Trees (BITs) of size
m+1, one for counts (bit_count) and one for sums (bit_sum). Also, initialize theanswerarray. - Process in Batches: Iterate through the sorted tuples. Process all tuples with the same
nums1value in a batch.- Query: For each tuple in the current batch:
- Get the total count of elements added to the BITs so far.
- If the total count is
kor less, the answer is the total sum frombit_sum. - Otherwise, perform a binary search on the ranks (
1tom) to find the rank of thek-th largest element (r_kth). This involves queryingbit_countinside the binary search. - Calculate the sum of the top
kelements usingr_kthand queries to both BITs. - Store the result in the
answerarray at the original index.
- Update: After querying for all tuples in the batch, update the BITs with the
nums2values from this batch.
- Query: For each tuple in the current batch:
Walkthrough
The core idea is to change the order of processing. Instead of calculating the answer for i=0, 1, ..., n-1 in order, we calculate it for indices sorted by their nums1 values. This way, as we move from an element i to j with nums1[i] < nums1[j], the set of candidates for j is a superset of candidates for i.
To efficiently manage the candidates and find the top k sum, we use Fenwick Trees on the ranks of nums2 values. The binary search to find the threshold rank r_kth works as follows: we are looking for the (totalCount - k + 1)-th smallest element. We binary search for a rank r and check if bit_count.query(r) is greater than or equal to our target count. This finds the threshold rank in O(log m) BIT queries, each taking O(log m) time.
import java.util.*; class Solution { public long[] getResults(int[] nums1, int[] nums2, int k) { int n = nums1.length; // Coordinate Compression Set<Integer> uniqueVals = new HashSet<>(); for (int val : nums2) uniqueVals.add(val); List<Integer> sortedUniqueVals = new ArrayList<>(uniqueVals); Collections.sort(sortedUniqueVals); Map<Integer, Integer> valToRank = new HashMap<>(); for (int i = 0; i < sortedUniqueVals.size(); i++) { valToRank.put(sortedUniqueVals.get(i), i + 1); } int m = sortedUniqueVals.size(); // Combine and Sort int[][] pairs = new int[n][3]; for (int i = 0; i < n; i++) { pairs[i] = new int[]{nums1[i], nums2[i], i}; } Arrays.sort(pairs, Comparator.comparingInt(p -> p[0])); BIT bit_count = new BIT(m + 1); BIT bit_sum = new BIT(m + 1); long[] answer = new long[n]; int i = 0; while (i < n) { int j = i; while (j < n && pairs[j][0] == pairs[i][0]) j++; for (int l = i; l < j; l++) { long totalCount = bit_count.query(m); if (totalCount < k) { answer[pairs[l][2]] = bit_sum.query(m); } else { long targetCount = totalCount - k + 1; int low = 1, high = m, r_kth = 0; while (low <= high) { int mid = low + (high - low) / 2; if (bit_count.query(mid) >= targetCount) { r_kth = mid; high = mid - 1; } else { low = mid + 1; } } long sum_gt_kth = bit_sum.query(m) - bit_sum.query(r_kth); long count_gt_kth = totalCount - bit_count.query(r_kth); long needed_at_kth = k - count_gt_kth; long sum_at_kth = needed_at_kth * sortedUniqueVals.get(r_kth - 1); answer[pairs[l][2]] = sum_gt_kth + sum_at_kth; } } for (int l = i; l < j; l++) { int rank = valToRank.get(pairs[l][1]); bit_count.update(rank, 1); bit_sum.update(rank, pairs[l][1]); } i = j; } return answer; }} class BIT { long[] tree; int size; BIT(int size) { this.size = size; this.tree = new long[size]; } void update(int index, long val) { while (index < size) { tree[index] += val; index += index & -index; } } long query(int index) { long sum = 0; while (index > 0) { sum += tree[index]; index -= index & -index; } return sum; }}Complexity
Time
O(n * (log n)^2) Sorting takes `O(n log n)`. The main loop processes `n` elements. Each query involves a binary search on ranks (`m <= n`), which takes `O(log m)`. Inside the binary search, we query the BIT, taking another `O(log m)`. This leads to `O((log m)^2)` per query. The total time is dominated by `n` such queries.
Space
O(n) We need `O(n)` space for the pairs, `O(m)` for coordinate compression data (where `m <= n`), and `O(m)` for each of the two Fenwick Trees.
Trade-offs
Pros
Much more efficient than the brute-force approach.
Passes the time limits for the given constraints.
Demonstrates a powerful technique for solving offline query problems.
Cons
Requires understanding of advanced data structures (Fenwick Tree).
Implementation is more complex than the brute-force approach.
The
O((log n)^2)factor per query can be slower than the fully optimized solution for some competitive programming scenarios.
This approach is a further optimization of the Fenwick Tree solution. It achieves the best possible time complexity by improving how we find the k-th largest element. The O((log m)^2) query time from the previous approach, caused by a binary search that calls a O(log m) BIT query, is reduced to just O(log m). This is done using a technique known as "walking on the BIT" or "binary lifting on the BIT". This specialized algorithm leverages the BIT's structure to perform a search for a specific cumulative frequency in logarithmic time, making the overall solution faster.
Algorithm
- The overall algorithm structure (coordinate compression, sorting, processing in batches) is identical to the previous approach.
- The key difference lies in the query step for finding the
k-th largest element's rank. - Optimized Query (BIT Walk): Instead of a
O((log m)^2)binary search, we use aO(log m)method called "walking on the BIT".- To find the rank of the
C-th smallest element, we traverse the BIT from top to bottom (from largest power of 2 downwards). - We maintain a position and try to extend it by the current power of two. If the count in the sub-tree does not exceed our target count, we move to that position and reduce the target count.
- This process efficiently pinpoints the desired rank
r_kthin logarithmic time.
- To find the rank of the
- Once
r_kthis found, the sum calculation is the same as in the previous approach.
Walkthrough
The logic remains the same as the O(n (log n)^2) solution, but we replace the binary search for r_kth with a more efficient findKth method. This method effectively performs a binary search on the prefix sums implicitly, without repeated calls to the query function. It navigates through the BIT's tree-like structure to find the index corresponding to a certain cumulative count.
// Add this method to the BIT class from the previous approach.class BIT { // ... existing fields and methods ... public int findKth(long k) { // finds the smallest index `i` such that query(i) >= k int pos = 0; long current_sum = 0; int log_size = 0; // Determine the highest power of 2 less than size for (int temp = size - 1; temp > 0; temp >>= 1) { log_size++; } for(int p = 1 << log_size; p > 0; p >>= 1){ if(pos + p < size && current_sum + tree[pos + p] < k){ current_sum += tree[pos + p]; pos += p; } } return pos + 1; }} // In the main logic, replace the binary search block with a single call:// ...long totalCount = bit_count.query(m);if (totalCount < k) { answer[pairs[l][2]] = bit_sum.query(m);} else { long targetCount = totalCount - k + 1; int r_kth = bit_count.findKth(targetCount); // The rest of the sum calculation is the same as before long sum_gt_kth = bit_sum.query(m) - bit_sum.query(r_kth); long count_gt_kth = totalCount - bit_count.query(r_kth); long needed_at_kth = k - count_gt_kth; long sum_at_kth = needed_at_kth * sortedUniqueVals.get(r_kth - 1); answer[pairs[l][2]] = sum_gt_kth + sum_at_kth;}// ...Complexity
Time
O(n log n) Sorting still takes `O(n log n)`. With the BIT walk optimization, each query to find the `k`-th element now takes `O(log m)`. Since `m <= n`, the total time for `n` queries and updates is `O(n log n)`. The overall complexity is therefore `O(n log n)`.
Space
O(n) The space requirements are identical to the previous Fenwick Tree approach.
Trade-offs
Pros
The most asymptotically efficient solution for this problem.
Achieves an optimal time complexity that is hard to beat.
Cons
The implementation of the BIT walk is non-trivial and requires a deeper understanding of the BIT's internal structure.
The code is more complex than the standard BIT approach.
Solutions
Solution
class Solution {public long[] findMaxSum(int[] nums1, int[] nums2, int k) { int n = nums1.length; int[][] arr = new int[n][0]; for (int i = 0; i < n; ++i) { arr[i] = new int[]{nums1[i], i}; } Arrays.sort(arr, (a, b)->a[0] - b[0]); PriorityQueue<Integer> pq = new PriorityQueue<>(); long s = 0; long[] ans = new long[n]; int j = 0; for (int h = 0; h < n; ++h) { int x = arr[h][0], i = arr[h][1]; while (j < h && arr[j][0] < x) { int y = nums2[arr[j][1]]; pq.offer(y); s += y; if (pq.size() > k) { s -= pq.poll(); } ++j; } ans[i] = s; } return ans; }}class Solution {public: vector<long long> findMaxSum(vector<int> &nums1, vector<int> &nums2, int k) { int n = nums1.size(); vector<pair<int, int>> arr(n); for (int i = 0; i < n; ++i) { arr[i] = {nums1[i], i}; } ranges ::sort(arr); priority_queue<int, vector<int>, greater<int>> pq; long long s = 0; int j = 0; vector<long long> ans(n); for (int h = 0; h < n; ++h) { auto [x, i] = arr[h]; while (j < h && arr[j].first < x) { int y = nums2[arr[j].second]; pq.push(y); s += y; if (pq.size() > k) { s -= pq.top(); pq.pop(); } ++j; } ans[i] = s; } return ans; }};class Solution: def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]: arr = [(x, i) for i, x in enumerate(nums1)] arr . sort() pq = [] s = j = 0 n = len(arr) ans = [0] * n for h, (x, i) in enumerate(arr): while j < h and arr[j][0] < x: y = nums2[arr[j][1]] heappush(pq, y) s += y if len(pq) > k: s -= heappop(pq) j += 1 ans[i] = s return ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.