Welcome to Subscribe On Youtube

3478. Choose K Elements With Maximum Sum

Description

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

Solutions

Solution 1: Sorting + Priority Queue (Min-Heap)

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

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

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

After traversing, we return the answer array $\mathit{\text{ans}}$.

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

• Java
• C++
• Python
• Go
• TypeScript

• 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 ans
  
  

• func findMaxSum(nums1 []int, nums2 []int, k int) []int64 {
  	n := len(nums1)
  	arr := make([][2]int, n)
  	for i, x := range nums1 {
  		arr[i] = [2]int{x, i}
  	}
  	ans := make([]int64, n)
  	sort.Slice(arr, func(i, j int) bool { return arr[i][0] < arr[j][0] })
  	pq := hp{}
  	var s int64
  	j := 0
  	for h, e := range arr {
  		x, i := e[0], e[1]
  		for j < h && arr[j][0] < x {
  			y := nums2[arr[j][1]]
  			heap.Push(&pq, y)
  			s += int64(y)
  			if pq.Len() > k {
  				s -= int64(heap.Pop(&pq).(int))
  			}
  			j++
  		}
  		ans[i] = s
  	}
  	return ans
  }
  
  type hp struct{ sort.IntSlice }
  
  func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] }
  func (h *hp) Push(v any)        { h.IntSlice = append(h.IntSlice, v.(int)) }
  func (h *hp) Pop() any {
  	a := h.IntSlice
  	v := a[len(a)-1]
  	h.IntSlice = a[:len(a)-1]
  	return v
  }
  
  

• function findMaxSum(nums1: number[], nums2: number[], k: number): number[] {
      const n = nums1.length;
      const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]);
      const pq = new MinPriorityQueue();
      let [s, j] = [0, 0];
      const ans: number[] = Array(k).fill(0);
      for (let h = 0; h < n; ++h) {
          const [x, i] = arr[h];
          while (j < h && arr[j][0] < x) {
              const y = nums2[arr[j++][1]];
              pq.enqueue(y);
              s += y;
              if (pq.size() > k) {
                  s -= pq.dequeue();
              }
          }
          ans[i] = s;
      }
      return ans;
  }
  
  

All Problems

All Solutions