You are given a 2D integer matrix grid of size n x m, an integer array limits of length n, and an integer k. The task is to find the maximum sum of at most k elements from the matrix grid such that:
-
The number of elements taken from the
ithrow ofgriddoes not exceedlimits[i].
Return the maximum sum.
Example 1:
Input: grid = [[1,2],[3,4]], limits = [1,2], k = 2
Output: 7
Explanation:
- From the second row, we can take at most 2 elements. The elements taken are 4 and 3.
- The maximum possible sum of at most 2 selected elements is
4 + 3 = 7.
Example 2:
Input: grid = [[5,3,7],[8,2,6]], limits = [2,2], k = 3
Output: 21
Explanation:
- From the first row, we can take at most 2 elements. The element taken is 7.
- From the second row, we can take at most 2 elements. The elements taken are 8 and 6.
- The maximum possible sum of at most 3 selected elements is
7 + 8 + 6 = 21.
Constraints:
n == grid.length == limits.lengthm == grid[i].length1 <= n, m <= 5000 <= grid[i][j] <= 1050 <= limits[i] <= m0 <= k <= min(n * m, sum(limits))
class Solution:
def maxSum(self, grid: List[List[int]], limits: List[int], k: int) -> int:
# Every row number i can't exceed limits[i]
# at most k
m, n = len(grid), len(grid[0])
pq = []
for row in range(m):
for col in range(n):
heappush(pq, (-grid[row][col], row))
ans = 0
rowsUsed = [0] * m
while k > 0:
k -= 1
while rowsUsed[pq[0][1]] >= limits[pq[0][1]]:
heappop(pq)
num, row = heappop(pq)
ans -= num
rowsUsed[row] += 1
return ans