Problem Description

You are given a m x n matrix mat containing integers and an integer value threshold. Your task is to find the maximum side-length of a square submatrix within mat such that the sum of all elements in that square is less than or equal to threshold.

The problem asks you to:

1. Consider all possible square submatrices within the given matrix
2. Calculate the sum of elements for each square
3. Find the largest square whose sum doesn't exceed the threshold
4. Return the side-length of this largest valid square

If no square exists with a sum less than or equal to threshold (including 1x1 squares), return 0.

For example, if you have a 3x3 matrix and threshold of 8, you would check:

• All 1x1 squares (individual elements)
• All 2x2 squares that can fit in the matrix
• All 3x3 squares (in this case, just the entire matrix)

The solution uses:

• Prefix sum array (s): A 2D prefix sum array where s[i][j] represents the sum of all elements in the rectangle from (0,0) to (i-1,j-1). This allows calculating any rectangular sum in O(1) time using the formula: sum = s[i+k][j+k] - s[i][j+k] - s[i+k][j] + s[i][j] for a k×k square starting at (i,j).

• Binary search: Since we're looking for the maximum side-length, binary search is used on the possible side-lengths (from 0 to min(m,n)). The check(k) function verifies if there exists at least one k×k square with sum ≤ threshold.

The algorithm efficiently finds the answer by combining prefix sums for quick sum calculations with binary search to minimize the number of side-lengths to check.

Intuition

The key insight is recognizing that this problem has a monotonic property: if a square of side-length k exists with sum ≤ threshold, then all squares with side-length less than k will also have valid squares (we can always find smaller squares within the larger valid square). Conversely, if no square of side-length k exists with sum ≤ threshold, then no larger squares will work either.

This monotonic property immediately suggests binary search on the answer. Instead of checking every possible side-length from 1 to min(m,n), we can use binary search to find the maximum valid side-length in O(log(min(m,n))) iterations.

The next challenge is: how do we efficiently check if a square of given side-length k exists with sum ≤ threshold?

A naive approach would calculate the sum for each k×k square by iterating through all its elements, taking O(k²) time per square. With potentially O(m*n) squares to check, this becomes too slow.

This is where 2D prefix sums come in. By preprocessing the matrix to build a prefix sum array, we can calculate the sum of any rectangular region in O(1) time. The formula for a rectangle from (i,j) to (i+k-1, j+k-1) becomes: sum = s[i+k][j+k] - s[i][j+k] - s[i+k][j] + s[i][j]

This works because:

• s[i+k][j+k] gives us the sum from origin to the bottom-right corner
• We subtract s[i][j+k] to remove the region above our square
• We subtract s[i+k][j] to remove the region to the left of our square
• We add back s[i][j] because we subtracted it twice in the previous two operations

Combining these two techniques - binary search for the side-length and prefix sums for quick validation - gives us an efficient solution that runs in O(m*n*log(min(m,n))) time.

Learn more about Binary Search and Prefix Sum patterns.

Solution Approach

The implementation consists of three main parts: building the prefix sum array, implementing the validation function, and performing binary search.

Step 1: Building the 2D Prefix Sum Array

We create a 2D array s with dimensions (m+1) × (n+1), padded with zeros to handle boundary cases elegantly:

s = [[0] * (n + 1) for _ in range(m + 1)]
for i, row in enumerate(mat, 1):
    for j, x in enumerate(row, 1):
        s[i][j] = s[i - 1][j] + s[i][j - 1] - s[i - 1][j - 1] + x

The formula s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + x builds the prefix sum incrementally:

• s[i-1][j]: sum of rectangle from origin to row i-1, column j
• s[i][j-1]: sum of rectangle from origin to row i, column j-1
• -s[i-1][j-1]: subtract the overlap (counted twice)
• +x: add the current element mat[i-1][j-1]

Step 2: Validation Function check(k)

This function determines if there exists at least one k×k square with sum ≤ threshold:

def check(k: int) -> bool:
    for i in range(m - k + 1):
        for j in range(n - k + 1):
            v = s[i + k][j + k] - s[i][j + k] - s[i + k][j] + s[i][j]
            if v <= threshold:
                return True
    return False

For each possible k×k square starting at position (i,j):

• We calculate its sum using the prefix sum array in O(1) time
• The range m - k + 1 ensures the square fits within matrix boundaries
• We return True immediately upon finding any valid square

Step 3: Binary Search for Maximum Side-Length

l, r = 0, min(m, n)
while l < r:
    mid = (l + r + 1) >> 1
    if check(mid):
        l = mid
    else:
        r = mid - 1
return l

The binary search works as follows:

• Initialize search range: l = 0 (minimum possible), r = min(m, n) (maximum possible square that fits)
• Use mid = (l + r + 1) >> 1 to avoid infinite loops when searching for the maximum value
• If check(mid) returns True, we know squares of size mid are valid, so search larger sizes: l = mid
• Otherwise, the current size is too large, search smaller sizes: r = mid - 1
• The loop terminates when l == r, giving us the maximum valid side-length

Time Complexity: O(m*n*log(min(m,n)))

• Building prefix sum: O(m*n)
• Binary search iterations: O(log(min(m,n)))
• Each check() call: O(m*n) to verify all possible positions

Space Complexity: O(m*n) for the prefix sum array

Example Walkthrough

Let's walk through a concrete example to illustrate the solution approach.

Given matrix:

mat = [[1, 1, 3],
       [2, 2, 3],
       [1, 1, 3]]
threshold = 4

Step 1: Build the 2D Prefix Sum Array

We create a prefix sum array s with dimensions 4×4 (padded with zeros):

s = [[0, 0, 0, 0],
     [0, 1, 2, 5],
     [0, 3, 6, 12],
     [0, 4, 8, 17]]

Let's trace how we calculate s[2][2]:

• s[2][2] = s[1][2] + s[2][1] - s[1][1] + mat[1][1]
• s[2][2] = 2 + 3 - 1 + 2 = 6

This represents the sum of all elements from (0,0) to (1,1) in the original matrix.

Step 2: Binary Search on Side-Length

Initial range: l = 0, r = min(3, 3) = 3

First iteration: mid = (0 + 3 + 1) >> 1 = 2

• Check if any 2×2 square has sum ≤ 4:
  • Square at (0,0): sum = s[2][2] - s[0][2] - s[2][0] + s[0][0] = 6 - 0 - 0 + 0 = 6 (> 4, invalid)
  • Square at (0,1): sum = s[2][3] - s[0][3] - s[2][1] + s[0][1] = 12 - 0 - 3 + 0 = 9 (> 4, invalid)
  • Square at (1,0): sum = s[3][2] - s[1][2] - s[3][0] + s[1][0] = 8 - 2 - 0 + 0 = 6 (> 4, invalid)
  • Square at (1,1): sum = s[3][3] - s[1][3] - s[3][1] + s[1][1] = 17 - 5 - 4 + 1 = 9 (> 4, invalid)
• check(2) returns False, so r = mid - 1 = 1

Second iteration: mid = (0 + 1 + 1) >> 1 = 1

• Check if any 1×1 square has sum ≤ 4:
  • Square at (0,0): sum = s[1][1] - s[0][1] - s[1][0] + s[0][0] = 1 - 0 - 0 + 0 = 1 (≤ 4, valid!)
• check(1) returns True, so l = mid = 1

Termination: l = r = 1, return 1

The maximum side-length of a square with sum ≤ 4 is 1.

Verification of the prefix sum formula: Let's verify the 2×2 square at position (0,0) manually:

• Elements: mat[0][0] + mat[0][1] + mat[1][0] + mat[1][1] = 1 + 1 + 2 + 2 = 6
• Using prefix sum: s[2][2] - s[0][2] - s[2][0] + s[0][0] = 6 - 0 - 0 + 0 = 6 ✓

The algorithm efficiently found that only 1×1 squares satisfy the threshold constraint, avoiding the need to check all possible combinations exhaustively.

Solution Implementation

Time and Space Complexity

Time Complexity: O(m * n * min(m, n))

The time complexity is determined by:

• Building the prefix sum array takes O(m * n) time
• Binary search runs for O(log(min(m, n))) iterations
• For each binary search iteration, the check function is called
• Inside check, there are two nested loops iterating (m - k + 1) * (n - k + 1) times, which in the worst case (when k is small) is approximately O(m * n)
• Each iteration inside check performs constant time operations using the prefix sum array

Therefore, the overall time complexity is O(m * n * log(min(m, n))). However, since the binary search explores different values of k, and for larger values of k the nested loops in check run fewer times, a tighter bound considering all iterations together gives us O(m * n * min(m, n)).

Space Complexity: O(m * n)

The space complexity is determined by:

• The prefix sum array s which has dimensions (m + 1) × (n + 1), requiring O(m * n) space
• Other variables like l, r, mid, and loop counters use O(1) space

Therefore, the overall space complexity is O(m * n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Binary Search Midpoint Calculation

Pitfall: Using mid = (left + right) // 2 when searching for the maximum value leads to an infinite loop.

When left = right - 1 and check(mid) returns True, setting left = mid keeps left at the same value, causing the loop to never terminate.

Example scenario:

• left = 2, right = 3
• mid = (2 + 3) // 2 = 2
• If check(2) returns True, we set left = 2
• Loop continues with same values indefinitely

Solution: Use mid = (left + right + 1) // 2 (or (left + right + 1) >> 1) to bias toward the upper value, ensuring progress in each iteration.

2. Off-by-One Errors in Prefix Sum Indexing

Pitfall: Confusing the coordinate system between the original matrix (0-indexed) and the prefix sum array (1-indexed with padding).

Common mistakes:

• Using prefix_sum[i][j] to represent sum up to mat[i][j] instead of mat[i-1][j-1]
• Incorrectly calculating square boundaries when checking positions

Solution: Maintain clear understanding that prefix_sum[i][j] represents the sum of all elements from mat[0][0] to mat[i-1][j-1]. When checking a k×k square starting at mat[i][j], use:

square_sum = (prefix_sum[i + k][j + k] 
             - prefix_sum[i][j + k] 
             - prefix_sum[i + k][j] 
             + prefix_sum[i][j])

3. Incorrect Range Calculation for Square Positions

Pitfall: Using range(m - k) or range(n - k) instead of range(m - k + 1) when iterating through possible square positions.

This misses valid squares that can start at the last valid positions. For example, in a 5×5 matrix, a 3×3 square can start at positions (0,0) through (2,2), requiring range(3) which is range(5 - 3 + 1).

Solution: Always use range(m - k + 1) and range(n - k + 1) to include all valid starting positions.

4. Not Handling Edge Cases

Pitfall: Failing to handle matrices where even a 1×1 square exceeds the threshold.

If all individual elements are greater than the threshold, the function should return 0, but incorrect initialization or boundary handling might cause errors.

Solution: Initialize the binary search with left = 0 and ensure the check function properly handles k = 0 or returns False when no valid squares exist.