跳转至

3462. 提取至多 K 个元素的最大总和

题目描述

给你一个大小为 n x m 的二维矩阵 grid ，以及一个长度为 n 的整数数组 limits ，和一个整数 k 。你的目标是从矩阵 grid 中提取出 至多 k 个元素，并计算这些元素的最大总和，提取时需满足以下限制：

  • 从 grid 的第 i 行提取的元素数量不超过 limits[i] 。

返回最大总和。

 

示例 1：

输入：grid = [[1,2],[3,4]], limits = [1,2], k = 2

输出：7

解释：

  • 从第 2 行提取至多 2 个元素，取出 4 和 3 。
  • 至多提取 2 个元素时的最大总和 4 + 3 = 7 。

示例 2：

输入：grid = [[5,3,7],[8,2,6]], limits = [2,2], k = 3

输出：21

解释：

  • 从第 1 行提取至多 2 个元素，取出 7 。
  • 从第 2 行提取至多 2 个元素，取出 8 和 6 。
  • 至多提取 3 个元素时的最大总和 7 + 8 + 6 = 21 。

 

提示：

  • n == grid.length == limits.length
  • m == grid[i].length
  • 1 <= n, m <= 500
  • 0 <= grid[i][j] <= 105
  • 0 <= limits[i] <= m
  • 0 <= k <= min(n * m, sum(limits))

解法

方法一：贪心 + 优先队列（小根堆）

我们可以用一个优先队列（小根堆） \(\textit{pq}\) 来维护最大的 \(k\) 个元素。

遍历每一行，对每一行的元素进行排序，然后取出每一行最大的 \(\textit{limit}\) 个元素，将这些元素加入 \(\textit{pq}\) 中。如果 \(\textit{pq}\) 的大小超过了 \(k\)，就将堆顶元素弹出。

最后，将 \(\textit{pq}\) 中的元素相加即可。

时间复杂度 \(O(n \times m \times (\log m + \log k))\)，空间复杂度 \(O(k)\)。其中 \(n\) 和 \(m\) 分别为矩阵 \(\textit{grid}\) 的行数和列数。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
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)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
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;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
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;
}

评论