3462. Maximum Sum With at Most K Elements 

Description

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 i^th row of grid does not exceed limits[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.length
m == grid[i].length
1 <= n, m <= 500
0 <= grid[i][j] <= 10⁵
0 <= limits[i] <= m
0 <= k <= min(n * m, sum(limits))

Solutions

Solution 1: Greedy + Priority Queue (Min-Heap)

We can use a priority queue (min-heap) $\textit{pq}$ to maintain the largest $k$ elements.

Traverse each row, sort the elements in each row, and then take the largest $\textit{limit}$ elements from each row and add them to $\textit{pq}$. If the size of $\textit{pq}$ exceeds $k$, pop the top element of the heap.

Finally, sum the elements in $\textit{pq}$.

The time complexity is $O(n \times m \times (\log m + \log k))$, and the space complexity is $O(k)$. Here, $n$ and $m$ are the number of rows and columns of the matrix $\textit{grid}$, respectively.

Python3

class Solution:
    def maxSum(self, grid: List[List[int]], limits: List[int], k: int) -> int:
        pq = []
        for nums, limit in zip(grid, limits):
            nums.sort()
            for _ in range(limit):
                heappush(pq, nums.pop())
                if len(pq) > k:
                    heappop(pq)
        return sum(pq)

Java

class Solution {
    public long maxSum(int[][] grid, int[] limits, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        int n = grid.length;
        for (int i = 0; i < n; ++i) {
            int[] nums = grid[i];
            int limit = limits[i];
            Arrays.sort(nums);
            for (int j = 0; j < limit; ++j) {
                pq.offer(nums[nums.length - j - 1]);
                if (pq.size() > k) {
                    pq.poll();
                }
            }
        }
        long ans = 0;
        for (int x : pq) {
            ans += x;
        }
        return ans;
    }
}

C++

class Solution {
public:
    long long maxSum(vector<vector<int>>& grid, vector<int>& limits, int k) {
        priority_queue<int, vector<int>, greater<int>> pq;
        int n = grid.size();

        for (int i = 0; i < n; ++i) {
            vector<int> nums = grid[i];
            int limit = limits[i];
            ranges::sort(nums);

            for (int j = 0; j < limit; ++j) {
                pq.push(nums[nums.size() - j - 1]);
                if (pq.size() > k) {
                    pq.pop();
                }
            }
        }

        long long ans = 0;
        while (!pq.empty()) {
            ans += pq.top();
            pq.pop();
        }

        return ans;
    }
};

Go

type MinHeap []int

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}
func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func maxSum(grid [][]int, limits []int, k int) int64 {
	pq := &MinHeap{}
	heap.Init(pq)
	n := len(grid)

	for i := 0; i < n; i++ {
		nums := make([]int, len(grid[i]))
		copy(nums, grid[i])
		limit := limits[i]
		sort.Ints(nums)

		for j := 0; j < limit; j++ {
			heap.Push(pq, nums[len(nums)-j-1])
			if pq.Len() > k {
				heap.Pop(pq)
			}
		}
	}

	var ans int64 = 0
	for pq.Len() > 0 {
		ans += int64(heap.Pop(pq).(int))
	}

	return ans
}

TypeScript

function maxSum(grid: number[][], limits: number[], k: number): number {
    const pq = new MinPriorityQueue();
    const n = grid.length;
    for (let i = 0; i < n; i++) {
        const nums = grid[i];
        const limit = limits[i];
        nums.sort((a, b) => a - b);
        for (let j = 0; j < limit; j++) {
            pq.enqueue(nums[nums.length - j - 1]);
            if (pq.size() > k) {
                pq.dequeue();
            }
        }
    }
    let ans = 0;
    while (!pq.isEmpty()) {
        ans += pq.dequeue() as number;
    }
    return ans;
}