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

给你一个大小为 m x n 的矩阵 mat 和一个整数阈值 threshold。

请你返回元素总和小于或等于阈值的正方形区域的最大边长；如果没有这样的正方形区域，则返回 0 。

示例 1：

输入：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

解法：二维前缀和 + 二分查找

我们首先计算出数组 mat 的前缀和数组 P，随后依次枚举 mat 中的正方形，计算出每个正方形的元素之和。

具体地，当数组 mat 的大小为 m * n 时，正方形的左上角可以是 mat 中的任意位置，边长不会超过 m 和 n 中的较小值 min(m, n)，这样我们就可以使用三重循环枚举所有的正方形，时间复杂度为 O(MN∗min(M,N))。

由于我们可以借助数组 P 在 O(1) 的时间计算任意正方形的元素之和，因此该算法的总时间复杂度为 O(MN∗min(M,N))。

由于数组 mat 中的所有元素均为非负整数，因此若存在一个边长为 c 且元素之和不超过阈值的正方形，那一定存在一个边长为 1, 2, ..., c - 1 且元素之和不超过阈值的正方形（在边长为 c 的正方形内任取一个边长为 1, 2, ..., c - 1 的正方形即可）。

这样我们可以使用二分查找的方法，找出最大的边长 c。二分查找的上界为 min(m, n)，下界为 1，在二分查找的过程中，若当前查找的边长为 c'，我们只需要枚举 mat 中所有边长为 c' 的正方形，并判断其中是否存在一个元素之和不超过阈值的正方形即可。

Java版：

class Solution {
    public int maxSideLength(int[][] mat, int threshold) {
        int m = mat.length;
        int n = mat[0].length;
        int[][] presum = new int[m + 1][n + 1];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                presum[i + 1][j + 1] = presum[i + 1][j] + presum[i][j + 1] - presum[i][j] + mat[i][j];
            }
        }
        int l = 1;
        int r = Math.min(m, n);
        int ans = 0;
        while (l <= r) {
            int mid = l + (r - l) / 2;
            boolean find = false;
            for (int i = 0; i <= m - mid; i++) {
                for (int j = 0; j <= n - mid; j++) {
                    if (getRect(presum, i, j, i + mid - 1, j + mid - 1) <= threshold) {
                        find = true;
                        break;
                    }
                }
                if (find) {
                    break;
                }
            }
            if (find) {
                l = mid + 1;
            } else {
                r = mid - 1;
            }
        }
        return r;
    }

    private int getRect(int[][] presum, int x1, int y1, int x2, int y2) {
        return presum[x2 + 1][y2 + 1] - presum[x2 + 1][y1] - presum[x1][y2 + 1] + presum[x1][y1];
    }
}

Python3版：

class Solution:
    def maxSideLength(self, mat: List[List[int]], threshold: int) -> int:
        m = len(mat)
        n = len(mat[0])
        presum = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m):
            for j in range(n):
                presum[i + 1][j + 1] = presum[i + 1][j] + presum[i][j + 1] - presum[i][j] + mat[i][j]
        
        def getRect(x1, y1, x2, y2):
            return presum[x2 + 1][y2 + 1] - presum[x1][y2 + 1] - presum[x2 + 1][y1] + presum[x1][y1]

        l = 1
        r = min(m, n)
        while l <= r:
            mid = l + (r - l) // 2
            find = False
            for i in range(m - mid + 1):
                for j in range(n - mid + 1):
                    if getRect(i, j, i + mid - 1, j + mid - 1) <= threshold:
                        find = True
                        break
                if find:
                    break
            if find:
                l = mid + 1
            else:
                r = mid - 1
        return r

复杂度分析

• 时间复杂度：O(MN∗log⁡min⁡(M,N))。二分查找的次数为 O(log⁡min⁡(M,N))，在每次二分查找中，需要枚举所有边长为 mid 的矩形，数量为 O(MN)，因此总时间复杂度为 O(MN∗log⁡min⁡(M,N))。
• 空间复杂度：O(MN)。