总和小于或等于K的最大尺寸正方形子矩阵

问题描述

给定一个m*n的矩阵，矩阵中各元素都是非负整数。现在定义一个子矩阵为原矩阵中的一些行和一些列所组成的矩阵。一个子矩阵的大小为其所包含的行数和列数的乘积。给定一个非负整数K，求原矩阵中的最大尺寸正方形子矩阵，使得子矩阵元素的总和不超过K。如果不存在这样的子矩阵，则返回0。

解决方案

思路

利用二分查找，找到最大的满足条件的正方形矩阵边长len，然后只需要判断是否存在一个子矩阵的边长为len，元素之和不超过K。利用前缀和，可以在$O(n^2)$的时间内求出以每个元素为右下角的子矩阵元素之和时的前缀和，进而可以在$O(n^3logM)$的时间复杂度内解决该问题。

代码示例

def maxSquareSubmatrix(matrix, k):
    """
    :type matrix: List[List[int]]
    :type k: int
    :rtype: int
    """
    m, n = len(matrix), len(matrix[0])

    # 边界条件判断，如果矩阵为空，返回0
    if m == 0 or n == 0:
        return 0

    # 初始化最大边长为最小值
    ans = float('-inf')
    left, right = 1, min(m, n)

    # 二分查找最大的满足条件的正方形矩阵边长len
    while left <= right:
        mid = left + (right - left) // 2
        if search(mid, matrix, k):
            ans = max(ans, mid)
            left = mid + 1
        else:
            right = mid - 1

    return ans ** 2 if ans != float('-inf') else 0

def search(len, matrix, k):
    """
    :type len: int
    :type matrix: List[List[int]]
    :type k: int
    :rtype: bool
    """
    m, n = len(matrix), len(matrix[0])

    # 初始化前缀和矩阵
    preSum = [[0] * (n + 1) for _ in range(m + 1)]

    for i in range(1, m + 1):
        for j in range(1, n + 1):
            preSum[i][j] = preSum[i - 1][j] + preSum[i][j - 1] - preSum[i - 1][j - 1] + matrix[i - 1][j - 1]

    # 枚举所有可能的正方形子矩阵，判断该子矩阵是否满足条件
    for i in range(len, m + 1):
        for j in range(len, n + 1):
            if preSum[i][j] - preSum[i - len][j] - preSum[i][j - len] + preSum[i - len][j - len] <= k:
                return True

    return False

复杂度分析

• 时间复杂度：$O(n^3logM)$
• 空间复杂度：$O(n^2)$

其中，$n$为矩阵的行列数，$M$为矩阵中元素值的最大值。