Skip to content

3462. Maximum Sum With at Most K Elements 👍

  • Time: $O(nm\log nm)$
  • Space: $O(k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
 public:
  long long maxSum(vector<vector<int>>& grid, vector<int>& limits, int k) {
    long ans = 0;
    priority_queue<int, vector<int>, greater<>> minHeap;

    for (int i = 0; i < grid.size(); ++i) {
      ranges::sort(grid[i], greater<>());
      for (int j = 0; j < limits[i]; ++j) {
        minHeap.push(grid[i][j]);
        if (minHeap.size() == k + 1)
          minHeap.pop();
      }
    }

    while (!minHeap.empty())
      ans += minHeap.top(), minHeap.pop();

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
  public long maxSum(int[][] grid, int[] limits, int k) {
    long ans = 0;
    Queue<Integer> minHeap = new PriorityQueue<>();

    for (int i = 0; i < grid.length; ++i) {
      Arrays.sort(grid[i]);
      for (int j = grid[i].length - 1; j >= grid[i].length - limits[i]; --j) {
        minHeap.offer(grid[i][j]);
        if (minHeap.size() == k + 1)
          minHeap.poll();
      }
    }

    while (!minHeap.isEmpty())
      ans += minHeap.poll();

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
class Solution:
  def maxSum(self, grid: list[list[int]], limits: list[int], k: int) -> int:
    minHeap = []

    for row, limit in zip(grid, limits):
      row.sort(reverse=True)
      for i in range(limit):
        heapq.heappush(minHeap, row[i])
        if len(minHeap) == k + 1:
          heapq.heappop(minHeap)

    return sum(minHeap)