Choose k elements with maximum sum

Array Sorting Heap (Priority Queue)

33.3% acceptance

Feb 25, 2026

178

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.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn find_max_sum(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<i64> {
    let n = nums1.len();
    let k = k as usize;
    // Sort by nums1 value. For each i, find all j where nums1[j] < nums1[i],
    // then take top k nums2[j] values.
    // Process in order of nums1: sort indices by nums1.
    let mut order: Vec<usize> = (0..n).collect();
    order.sort_unstable_by_key(|&i| nums1[i]);
    
    // Use a sorted structure to maintain top-k nums2 values seen so far.
    // A min-heap of size k: if heap.len() < k, push; else if top < new, pop and push.
    use std::collections::BinaryHeap;
    use std::cmp::Reverse;
    let mut heap: BinaryHeap<Reverse<i32>> = BinaryHeap::new(); // min-heap
    let mut sum_top_k = 0i64;
    let mut result = vec![0i64; n];
    
    let mut i = 0;
    while i < n {
      let val = nums1[order[i]];
      // Process all j with same nums1[j] value (they can't use each other)
      let mut j = i;
      while j < n && nums1[order[j]] == val { j += 1; }
      // For all indices in [i, j), result[order[x]] = sum_top_k
      for x in i..j { result[order[x]] = sum_top_k; }
      // Add nums2 values for indices [i, j) to heap
      for x in i..j {
        let v = nums2[order[x]];
        if heap.len() < k {
          heap.push(Reverse(v));
          sum_top_k += v as i64;
        } else if let Some(&Reverse(top)) = heap.peek() {
          if v > top {
            heap.pop(); sum_top_k -= top as i64;
            heap.push(Reverse(v)); sum_top_k += v as i64;
          }
        }
      }
      i = j;
    }
    result
  }
}