Problem Statement in English

You’re given an m x n matrix mat and an integer threshold. Return the maximum side length of a square with a sum less than or equal to threshold or return 0 if there is no such square.

Approach

The idea is to be able to find the sum of any submatrix in constant time. Then we can just use binary search on the possible side lengths of the square to find the maximum side length that meets the criteria efficiently.

To do this, we can use a 2D prefix sum. Once we have the prefix sums, we can use binary search to find the maximum side length of a square that meets the criteria.

The way we find a prefix sum is outlined in this problem I’ve done before.

Hook that up with binary search to find the maximum side length.

And we’re done.

Solution in Python


class Solution:
    def maxSideLength(self, matrix: List[List[int]], k: int) -> int:
        ROWS, COLS = len(matrix) + 1, len(matrix[0]) + 1

        backing = [[0] * COLS for _ in range(ROWS)]

        for i in range(1, ROWS):
            rowSum = 0
            for j in range(1, COLS):
                rowSum += matrix[i - 1][j - 1]
                backing[i][j] = rowSum + backing[i - 1][j]

        def good(side):
            for i in range(side, ROWS):
                # bottom right row
                for j in range(side, COLS):
                    # bottom right col

                    upper = backing[i - side][j]
                    left = backing[i][j - side]
                    overlapping = backing[i - side][j - side]
                    this = backing[i][j]

                    total = this - upper - left + overlapping
                    if total <= k:
                        return True
            
            return False

        l, r = 0, min(ROWS - 1, COLS - 1)

        while l <= r:
            m = (l + r) // 2

            if good(m):
                l = m + 1
            else:
                r = m - 1
        
        return r

Complexity

  • Time: $O((M \times N) \times \log(\min(M, N)))$
    Since we use binary search over the possible side lengths (up to $\min(M, N)$) and for each side length we check all possible squares in the matrix which takes $O(M \times N)$ time.

  • Space: $O(M \times N)$
    Since we are using a backing matrix of size $M \times N$ where $M$ and $N$ are the dimensions of the input matrix.

Mistakes I Made

I had to revise 2D prefix sums to implement this solution effectively. I also made a mistake with the binary search boundaries initially.

And we are done.