1292. Maximum Side Length of a Square with Sum Less than or Equal to Threshold

MediumArrayBinary SearchMatrixPrefix Sum

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:

Consider all possible square submatrices within the given matrix
Calculate the sum of elements for each square
Find the largest square whose sum doesn't exceed the threshold
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.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

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 Implementation

1class Solution:
2    def maxSideLength(self, mat: List[List[int]], threshold: int) -> int:
3        """
4        Find the maximum side length of a square with sum <= threshold.
5        Uses binary search on the side length and prefix sum for efficient sum calculation.
6
7        Args:
8            mat: 2D matrix of integers
9            threshold: Maximum allowed sum for a square submatrix
10
11        Returns:
12            Maximum side length of a valid square
13        """
14
15        def can_form_square_with_side(side_length: int) -> bool:
16            """
17            Check if there exists a square of given side length with sum <= threshold.
18
19            Args:
20                side_length: The side length of square to check
21
22            Returns:
23                True if such a square exists, False otherwise
24            """
25            # Try all possible top-left corners for a square of given side length
26            for row in range(rows - side_length + 1):
27                for col in range(cols - side_length + 1):
28                    # Calculate sum using prefix sum array
29                    # Sum of square from (row, col) to (row + side_length - 1, col + side_length - 1)
30                    square_sum = (prefix_sum[row + side_length][col + side_length]
31                                 - prefix_sum[row][col + side_length]
32                                 - prefix_sum[row + side_length][col]
33                                 + prefix_sum[row][col])
34
35                    if square_sum <= threshold:
36                        return True
37            return False
38
39        # Get matrix dimensions
40        rows, cols = len(mat), len(mat[0])
41
42        # Build prefix sum array with padding for easier calculation
43        # prefix_sum[i][j] represents sum of rectangle from (0,0) to (i-1, j-1)
44        prefix_sum = [[0] * (cols + 1) for _ in range(rows + 1)]
45
46        # Populate prefix sum array
47        for i, row in enumerate(mat, start=1):
48            for j, value in enumerate(row, start=1):
49                prefix_sum[i][j] = (prefix_sum[i - 1][j]
50                                   + prefix_sum[i][j - 1]
51                                   - prefix_sum[i - 1][j - 1]
52                                   + value)
53
54        # Binary search on the side length
55        # Maximum possible side length is the minimum of rows and cols
56        left, right = 0, min(rows, cols)
57
58        while left < right:
59            # Use ceiling division to avoid infinite loop
60            mid = (left + right + 1) >> 1  # Equivalent to (left + right + 1) // 2
61
62            if can_form_square_with_side(mid):
63                # If we can form a square with side length mid, try larger
64                left = mid
65            else:
66                # If we cannot form a square with side length mid, try smaller
67                right = mid - 1
68
69        return left
70

1class Solution {
2    // Matrix dimensions
3    private int rows;
4    private int cols;
5    private int threshold;
6    // Prefix sum array for calculating submatrix sums efficiently
7    private int[][] prefixSum;
8
9    public int maxSideLength(int[][] mat, int threshold) {
10        // Initialize dimensions
11        rows = mat.length;
12        cols = mat[0].length;
13        this.threshold = threshold;
14
15        // Build prefix sum array with extra row and column for easier calculation
16        // prefixSum[i][j] represents sum of all elements from (0,0) to (i-1,j-1)
17        prefixSum = new int[rows + 1][cols + 1];
18        for (int i = 1; i <= rows; i++) {
19            for (int j = 1; j <= cols; j++) {
20                // Calculate prefix sum using inclusion-exclusion principle
21                prefixSum[i][j] = prefixSum[i - 1][j]           // Sum above
22                                + prefixSum[i][j - 1]           // Sum to the left
23                                - prefixSum[i - 1][j - 1]       // Remove double counted area
24                                + mat[i - 1][j - 1];            // Add current element
25            }
26        }
27
28        // Binary search for the maximum side length
29        int left = 0;
30        int right = Math.min(rows, cols);  // Maximum possible side length
31
32        while (left < right) {
33            // Use ceiling division to avoid infinite loop
34            int mid = (left + right + 1) >> 1;  // Equivalent to (left + right + 1) / 2
35
36            if (check(mid)) {
37                // If a square of size mid exists, search for larger squares
38                left = mid;
39            } else {
40                // If no square of size mid exists, search for smaller squares
41                right = mid - 1;
42            }
43        }
44
45        return left;
46    }
47
48    /**
49     * Check if there exists a square submatrix with side length k
50     * whose sum is less than or equal to threshold
51     *
52     * @param k the side length to check
53     * @return true if such a square exists, false otherwise
54     */
55    private boolean check(int k) {
56        // Iterate through all possible top-left corners of k×k squares
57        for (int i = 0; i < rows - k + 1; i++) {
58            for (int j = 0; j < cols - k + 1; j++) {
59                // Calculate sum of k×k square starting at (i, j) using prefix sums
60                // Bottom-right corner is at (i+k-1, j+k-1)
61                int squareSum = prefixSum[i + k][j + k]    // Total sum to bottom-right
62                              - prefixSum[i][j + k]         // Subtract upper portion
63                              - prefixSum[i + k][j]         // Subtract left portion
64                              + prefixSum[i][j];            // Add back double-subtracted area
65
66                // If we find any valid square, return true
67                if (squareSum <= threshold) {
68                    return true;
69                }
70            }
71        }
72        return false;
73    }
74}
75

1class Solution {
2public:
3    int maxSideLength(vector<vector<int>>& mat, int threshold) {
4        int rows = mat.size();
5        int cols = mat[0].size();
6
7        // Create prefix sum array with padding for easier calculation
8        // prefixSum[i][j] represents the sum of all elements from (0,0) to (i-1,j-1)
9        vector<vector<int>> prefixSum(rows + 1, vector<int>(cols + 1, 0));
10
11        // Build the prefix sum array
12        // For each cell, add current element and use inclusion-exclusion principle
13        for (int i = 1; i <= rows; ++i) {
14            for (int j = 1; j <= cols; ++j) {
15                prefixSum[i][j] = prefixSum[i - 1][j] + prefixSum[i][j - 1]
16                                 - prefixSum[i - 1][j - 1] + mat[i - 1][j - 1];
17            }
18        }
19
20        // Lambda function to check if a square of side length k exists with sum <= threshold
21        auto canFormSquare = [&](int sideLength) -> bool {
22            // Try all possible top-left corners for a square of given side length
23            for (int i = 0; i <= rows - sideLength; ++i) {
24                for (int j = 0; j <= cols - sideLength; ++j) {
25                    // Calculate sum of square using prefix sum array
26                    // Bottom-right corner: (i + sideLength, j + sideLength)
27                    // Top-left corner: (i, j)
28                    int squareSum = prefixSum[i + sideLength][j + sideLength]
29                                  - prefixSum[i][j + sideLength]
30                                  - prefixSum[i + sideLength][j]
31                                  + prefixSum[i][j];
32
33                    // If we find any valid square, return true
34                    if (squareSum <= threshold) {
35                        return true;
36                    }
37                }
38            }
39            return false;
40        };
41
42        // Binary search for the maximum side length
43        int left = 0;
44        int right = min(rows, cols);
45
46        while (left < right) {
47            // Use ceiling division to avoid infinite loop
48            int mid = left + (right - left + 1) / 2;
49
50            if (canFormSquare(mid)) {
51                // If current side length works, try larger values
52                left = mid;
53            } else {
54                // If current side length doesn't work, try smaller values
55                right = mid - 1;
56            }
57        }
58
59        return left;
60    }
61};
62

1/**
2 * Finds the maximum side length of a square with sum less than or equal to threshold
3 * @param mat - The input matrix
4 * @param threshold - The maximum allowed sum for a square
5 * @returns The maximum side length of a valid square
6 */
7function maxSideLength(mat: number[][], threshold: number): number {
8    const rows: number = mat.length;
9    const cols: number = mat[0].length;
10
11    // Build prefix sum array for efficient range sum queries
12    // prefixSum[i][j] represents sum of all elements from (0,0) to (i-1,j-1)
13    const prefixSum: number[][] = Array(rows + 1)
14        .fill(0)
15        .map(() => Array(cols + 1).fill(0));
16
17    // Calculate prefix sums using inclusion-exclusion principle
18    for (let row = 1; row <= rows; row++) {
19        for (let col = 1; col <= cols; col++) {
20            prefixSum[row][col] = prefixSum[row - 1][col]
21                                 + prefixSum[row][col - 1]
22                                 - prefixSum[row - 1][col - 1]
23                                 + mat[row - 1][col - 1];
24        }
25    }
26
27    /**
28     * Checks if there exists a square of given side length with sum <= threshold
29     * @param sideLength - The side length to check
30     * @returns True if such a square exists, false otherwise
31     */
32    const hasValidSquare = (sideLength: number): boolean => {
33        // Try all possible top-left corners for a square of given side length
34        for (let topRow = 0; topRow <= rows - sideLength; topRow++) {
35            for (let leftCol = 0; leftCol <= cols - sideLength; leftCol++) {
36                // Calculate sum of square using prefix sums
37                // Bottom-right corner coordinates in prefix sum array
38                const bottomRow: number = topRow + sideLength;
39                const rightCol: number = leftCol + sideLength;
40
41                // Apply inclusion-exclusion to get square sum
42                const squareSum: number = prefixSum[bottomRow][rightCol]
43                                        - prefixSum[bottomRow][leftCol]
44                                        - prefixSum[topRow][rightCol]
45                                        + prefixSum[topRow][leftCol];
46
47                if (squareSum <= threshold) {
48                    return true;
49                }
50            }
51        }
52        return false;
53    };
54
55    // Binary search for maximum valid side length
56    let left: number = 0;
57    let right: number = Math.min(rows, cols);
58
59    while (left < right) {
60        // Use ceiling division to avoid infinite loop
61        const mid: number = (left + right + 1) >> 1;
62
63        if (hasValidSquare(mid)) {
64            // If mid is valid, search for larger side lengths
65            left = mid;
66        } else {
67            // If mid is invalid, search for smaller side lengths
68            right = mid - 1;
69        }
70    }
71
72    return left;
73}
74

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

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

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

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

How does merge sort divide the problem into subproblems?

Divide the array into two equal halves by index

Divide the array into two based on whether an element is smaller than an arbitrary value

Divide the array into a stray element and the rest of the array

Merge sort does not use divide and conquer

I don't know