# Choose K Elements With Maximum Sum
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/choose-k-elements-with-maximum-sum)
Canonical: https://scaleengineer.com/dsa/problems/choose-k-elements-with-maximum-sum
**Algorithms:** [Sorting](https://scaleengineer.com/algorithms/sorting)
**Data structures:** Array, Heap (Priority Queue)
---
## 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`.

**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`

# Approaches
## Brute Force Approach
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]`.
**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.
**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.
### Explanation
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`.

```java
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;
    }
}
```
### Algorithm
*   Initialize an empty `answer` array of size `n`.
*   Iterate through each index `i` from `0` to `n-1`.
    *   For each `i`, create a temporary list to store candidate `nums2` values.
    *   Iterate through all indices `j` from `0` to `n-1`.
    *   If `nums1[j] < nums1[i]`, add `nums2[j]` to the list of candidates.
    *   To find the sum of the top `k` candidates, use a min-priority queue of size `k`.
    *   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]`.
*   Return the `answer` array.

## Sorting with Fenwick Tree
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.
**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.
**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.
### Explanation
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.

```java
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;
    }
}
```
### Algorithm
*   **Coordinate Compression**: Create a sorted list of unique values from `nums2` and a map to get the rank of each value. Let `m` be the number of unique values.
*   **Combine and Sort**: Create tuples of `(nums1[i], nums2[i], original_index_i)` and sort these tuples based on `nums1` in 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 the `answer` array.
*   **Process in Batches**: Iterate through the sorted tuples. Process all tuples with the same `nums1` value 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 `k` or less, the answer is the total sum from `bit_sum`.
        *   Otherwise, perform a binary search on the ranks (`1` to `m`) to find the rank of the `k`-th largest element (`r_kth`). This involves querying `bit_count` inside the binary search.
        *   Calculate the sum of the top `k` elements using `r_kth` and queries to both BITs.
        *   Store the result in the `answer` array at the original index.
    *   **Update**: After querying for all tuples in the batch, update the BITs with the `nums2` values from this batch.

## Optimal Solution with Fenwick Tree and Walk Optimization
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.
**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.
**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.
### Explanation
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.

```java
// 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;
}
// ...
```
### 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 a `O(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_kth` in logarithmic time.
*   Once `r_kth` is found, the sum calculation is the same as in the previous approach.

# Solutions
### Java

```java
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;
  }
}

```

### CPP

```cpp
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;
  }
};

```

### Python

```python
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 ans

```