Problem Description

You are given a 2D matrix of integers called mat with dimensions m x n, and an integer threshold. Your task is to find the maximum possible side length of a square (with sides parallel to the axes) such that the sum of all elements inside the square is less than or equal to threshold.

• You can choose any square within the matrix, as long as all of its elements fit within the boundaries of mat.
• Each square must have all its elements' sum ≤ threshold.
• Return the largest possible side length (an integer) of such a square. If no such square exists, return 0.

Constraints:

• 1 ≤ m, n ≤ 300
• 0 ≤ mat[i][j] ≤ 10,000
• 0 ≤ threshold ≤ 10^5

Thought Process

At first glance, a brute-force approach would be to check every possible square in the matrix, for every possible size, and compute the sum of its elements. However, this would be extremely slow for large matrices, since there are many possible squares and each sum could take O(k²) time to compute for a square of size k.

To optimize, we need a way to quickly compute the sum of any sub-square in constant time. This leads us to the concept of a prefix sum matrix, which allows us to get the sum of any rectangular submatrix in O(1) time after preprocessing. With this, we can efficiently check all possible squares for a given side length.

To further optimize, instead of checking all possible square sizes one by one, we can use binary search to find the largest valid side length. This reduces the number of checks we need to make.

Solution Approach

The solution involves two main optimizations: prefix sum for fast area queries, and binary search for the answer.

1. Compute the Prefix Sum Matrix:
   • Create a matrix prefix of size (m+1) x (n+1), where prefix[i][j] is the sum of elements in the submatrix from (0,0) to (i-1,j-1).
   • This allows us to compute the sum of any submatrix from (r1,c1) to (r2,c2) in O(1) time using the inclusion-exclusion principle:
     sum = prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1]
2. Binary Search on Side Length:
   • Set your search range for side length from 0 up to min(m, n).
   • For each candidate side length k, check if there exists a square of size k x k whose sum is ≤ threshold.
   • If such a square exists, try a larger k; otherwise, try a smaller k.
3. Checking Squares Efficiently:
   • For a given side length k, slide a k x k window over the matrix.
   • Use the prefix sum matrix to compute the sum of each window in O(1) time.
   • If any such window's sum is ≤ threshold, the side length is possible.
4. Return the Largest Valid Side Length:
   • The binary search ensures we find the maximum k for which such a square exists.

Example Walkthrough

Consider mat = [[1,1,3,2,4,3,2],[1,1,3,2,4,3,2],[1,1,3,2,4,3,2]] and threshold = 4.

1. Build the prefix sum matrix.
2. Binary search for side length:
   • Try k=1: Every 1x1 square is just a single element, all are ≤ 4, so k=1 is possible.
   • Try k=2: For each 2x2 square, check if sum ≤ 4. For top-left square (positions (0,0)-(1,1)), sum is 1+1+1+1=4, which is valid. So k=2 is possible.
   • Try k=3: For top-left 3x3 square, sum is much larger (since there are 9 elements, many are >1), so the sum will exceed 4. So k=3 is not possible.
3. Conclusion: The answer is 2.

Time and Space Complexity

• Brute-force approach:
  • For each possible square (O(m * n * min(m, n))), check sum by iterating over all its elements (O(k^2)).
  • Total: O(m * n * (min(m, n))^3), which is too slow for large matrices.
• Optimized approach (Prefix Sum + Binary Search):
  • Prefix sum matrix construction: O(m * n) time and space.
  • Binary search over possible side lengths: O(log(min(m, n))).
  • For each binary search iteration, we check all possible k x k squares: O(m * n) per check.
  • Total: O(m * n * log(min(m, n))) time, O(m * n) space.

Summary

By combining prefix sums for fast area queries and binary search for the answer, we efficiently solve the problem of finding the largest square with sum ≤ threshold. The key insight is to preprocess the matrix for O(1) sum queries and to avoid checking every possible square size sequentially. This approach is both elegant and highly efficient, making it suitable for large input sizes.

Code Implementation

class Solution:
    def maxSideLength(self, mat, threshold):
        m, n = len(mat), len(mat[0])
        # Build prefix sum matrix
        prefix = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m):
            for j in range(n):
                prefix[i+1][j+1] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j]
        
        def possible(k):
            for i in range(m - k + 1):
                for j in range(n - k + 1):
                    total = prefix[i+k][j+k] - prefix[i][j+k] - prefix[i+k][j] + prefix[i][j]
                    if total <= threshold:
                        return True
            return False
        
        left, right = 0, min(m, n)
        ans = 0
        while left <= right:
            mid = (left + right) // 2
            if possible(mid):
                ans = mid
                left = mid + 1
            else:
                right = mid - 1
        return ans
    

class Solution {
public:
    int maxSideLength(vector<vector<int>>& mat, int threshold) {
        int m = mat.size(), n = mat[0].size();
        vector<vector<int>> prefix(m+1, vector<int>(n+1, 0));
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                prefix[i+1][j+1] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j];
            }
        }
        auto possible = [&](int k) {
            for (int i = 0; i + k <= m; ++i) {
                for (int j = 0; j + k <= n; ++j) {
                    int total = prefix[i+k][j+k] - prefix[i][j+k] - prefix[i+k][j] + prefix[i][j];
                    if (total <= threshold)
                        return true;
                }
            }
            return false;
        };
        int left = 0, right = min(m, n), ans = 0;
        while (left <= right) {
            int mid = (left + right) / 2;
            if (possible(mid)) {
                ans = mid;
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        return ans;
    }
};
    

class Solution {
    public int maxSideLength(int[][] mat, int threshold) {
        int m = mat.length, n = mat[0].length;
        int[][] prefix = new int[m + 1][n + 1];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                prefix[i + 1][j + 1] = mat[i][j] + prefix[i][j + 1] + prefix[i + 1][j] - prefix[i][j];
            }
        }
        int left = 0, right = Math.min(m, n), ans = 0;
        while (left <= right) {
            int mid = (left + right) / 2;
            if (possible(prefix, m, n, mid, threshold)) {
                ans = mid;
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        return ans;
    }
    private boolean possible(int[][] prefix, int m, int n, int k, int threshold) {
        for (int i = 0; i + k <= m; i++) {
            for (int j = 0; j + k <= n; j++) {
                int total = prefix[i + k][j + k] - prefix[i][j + k] - prefix[i + k][j] + prefix[i][j];
                if (total <= threshold)
                    return true;
            }
        }
        return false;
    }
}
    

var maxSideLength = function(mat, threshold) {
    const m = mat.length, n = mat[0].length;
    const prefix = Array.from({length: m + 1}, () => Array(n + 1).fill(0));
    for (let i = 0; i < m; i++) {
        for (let j = 0; j < n; j++) {
            prefix[i+1][j+1] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j];
        }
    }
    function possible(k) {
        for (let i = 0; i + k <= m; i++) {
            for (let j = 0; j + k <= n; j++) {
                let total = prefix[i+k][j+k] - prefix[i][j+k] - prefix[i+k][j] + prefix[i][j];
                if (total <= threshold) return true;
            }
        }
        return false;
    }
    let left = 0, right = Math.min(m, n), ans = 0;
    while (left <= right) {
        let mid = Math.floor((left + right) / 2);
        if (possible(mid)) {
            ans = mid;
            left = mid + 1;
        } else {
            right = mid - 1;
        }
    }
    return ans;
};