1292. 元素和小于等于阈值的正方形的最大边长

Difficulty
Medium
Tags
二分查找
前缀和
URL
https://leetcode-cn.com/problems/maximum-side-length-of-a-square-with-sum-less-than-or-equal-to-threshold/
Star
给你一个大小为 m x n 的矩阵 mat 和一个整数阈值 threshold。
请你返回元素总和小于或等于阈值的正方形区域的最大边长；如果没有这样的正方形区域，则返回 0 。
示例 1：
notion imagenotion image
输入：mat = [[1,1,3,2,4,3,2],[1,1,3,2,4,3,2],[1,1,3,2,4,3,2]], threshold = 4
输出：2
解释：总和小于或等于 4 的正方形的最大边长为 2，如图所示。
示例 2：
输入：mat = [[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2]], threshold = 1
输出：0
示例 3：
输入：mat = [[1,1,1,1],[1,0,0,0],[1,0,0,0],[1,0,0,0]], threshold = 6
输出：3
示例 4：
输入：mat = [[18,70],[61,1],[25,85],[14,40],[11,96],[97,96],[63,45]], threshold = 40184
输出：2
提示：
  • 1 <= m, n <= 300
  • m == mat.length
  • n == mat[i].length
  • 0 <= mat[i][j] <= 10000
  • 0 <= threshold <= 10^5

法 1 二分 + 前缀

思路
采用二分法来猜测答案。
具体判断的函数，用前缀和数组优化。
题解
class Solution:
    def maxSideLength(self, mat: List[List[int]], threshold: int) -> int:
        # 计算矩阵前缀和
        pre_sum = [[0] * (len(mat[0])+1) for _ in range(len(mat)+1)]
        for i in range(1, len(pre_sum)):
            for j in range(1, len(pre_sum[0])):
                pre_sum[i][j] = pre_sum[i-1][j] + pre_sum[i][j-1] - pre_sum[i-1][j-1] + mat[i-1][j-1]
        
        # 二分法查找答案
        left, right = 1, min(len(mat), len(mat[0]))
        # 找最大边长，所以是找右边界
        while left <= right:
            mid = left + (right - left) // 2
            if self.isValid(pre_sum, threshold, mid):
                left = mid + 1
            else:
                right = mid - 1
            
        return right
        
    def isValid(self, pre_sum, threshold, mid):
        for i in range(mid, len(pre_sum)):
            for j in range(mid, len(pre_sum[0])):
                # 计算每个边长为 mid 的正方形的 元素和
                if pre_sum[i][j] + pre_sum[i-mid][j-mid] - pre_sum[i-mid][j] - pre_sum[i][j-mid] <= threshold:
                    return True 
        return False