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.

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. Using Wrong Binary Search Template Variant

This is a "find maximum" problem, not "find minimum". The template must be adapted.

Wrong approach (causes infinite loop):

while left < right:
    mid = (left + right) // 2
    if feasible(mid):
        left = mid  # Infinite loop when left = right - 1
    else:
        right = mid - 1

Correct approach using template:

# Reframe: find first infeasible size, answer = first_infeasible - 1
first_true_index = max_possible + 1
while left <= right:
    mid = (left + right) // 2
    if not feasible(mid):  # Infeasible is our "true"
        first_true_index = mid
        right = mid - 1
    else:
        left = mid + 1
return first_true_index - 1

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).

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) instead of range(m - k + 1) when iterating through possible square positions.

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.

Solution: By searching for first infeasible and returning first_true_index - 1, we naturally handle this case. If all sizes are infeasible (first_true_index = 1), we return 0.