Solutions CS101

Appearance

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

data structures, https://leetcode.cn/problems/maximum-sum-with-at-most-k-elements/

给你一个大小为 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] <= 10^5
  • 0 <= limits[i] <= m
  • 0 <= k <= min(n * m, sum(limits))
python
from typing import List
import heapq

class Solution:
    def maxSum(self, grid: List[List[int]], limits: List[int], k: int) -> int:
        hqs = []
        n = len(grid)
        for i in range(n):
            row = grid[i]
            limit = limits[i]
            largest_k = sorted(row, reverse=True)
            if limit < len(largest_k):
                largest_k = largest_k[:limit]
            hq = [-val for val in largest_k]
            heapq.heapify(hq)

            if hq:
                hqs.append(hq)

        print(hqs)
        sum_v = 0
        for i in range(k):
            max_v = float('-inf')
            max_idx = -1

            for j, hq in enumerate(hqs):
                if hq and -hq[0] > max_v:
                    max_v = -hq[0]
                    max_idx = j

            if max_idx != -1:
                sum_v += max_v
                heapq.heappop(hqs[max_idx])
                if not hqs[max_idx]:
                    del hqs[max_idx]

        return sum_v

if __name__ == "__main__":
    sol = Solution()
    #print(sol.maxSum([[1,2],[3,4]], [1,2], 2))
    #print(sol.maxSum([[5,3,7],[8,2,6]], [2,2], 3))
    #print(sol.maxSum([[3],[9],[1]], [1,0,0], 1))
    print(sol.maxSum([[5,3,7],[8,2,6]], [2,2], 3))