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

medium prefix sum matrix binary search

Problem

Given an m × n matrix mat and an integer threshold, return the maximum side length of a square sub-grid whose element sum is less than or equal to threshold. If no such square exists, return 0.

The trick is a 2D prefix-sum table: pre[i][j] holds the sum of every element above-and-left of cell (i, j), so the sum of any axis-aligned rectangle (and thus any square) can be read in O(1).

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

Output2

A 2×2 square of two 1s and two 1s in the top-left sums to 4 ≤ 4. No 3×3 square stays within 4, so the answer is 2.

mat (rows separated by ';', values by ',')

threshold

def maxSideLength(mat, threshold):
    m, n = len(mat), len(mat[0])
    # pre[i][j] = sum of mat[0..i-1][0..j-1]
    pre = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            pre[i][j] = (mat[i - 1][j - 1] + pre[i - 1][j]
                         + pre[i][j - 1] - pre[i - 1][j - 1])

    def square_sum(r, c, k):           # square top-left (r, c), side k
        return (pre[r + k][c + k] - pre[r][c + k]
                - pre[r + k][c] + pre[r][c])

    best = 0
    for r in range(m):                 # top-left row
        for c in range(n):             # top-left column
            # greedily grow while a bigger square still fits and is light enough
            while best + 1 <= min(m - r, n - c) and \
                  square_sum(r, c, best + 1) <= threshold:
                best += 1
    return best

function maxSideLength(mat, threshold) {
  const m = mat.length, n = mat[0].length;
  // pre[i][j] = sum of mat[0..i-1][0..j-1]
  const pre = Array.from({ length: m + 1 }, () => new Array(n + 1).fill(0));
  for (let i = 1; i <= m; i++)
    for (let j = 1; j <= n; j++)
      pre[i][j] = mat[i - 1][j - 1] + pre[i - 1][j]
                + pre[i][j - 1] - pre[i - 1][j - 1];

  const squareSum = (r, c, k) =>        // square top-left (r, c), side k
    pre[r + k][c + k] - pre[r][c + k] - pre[r + k][c] + pre[r][c];

  let best = 0;
  for (let r = 0; r < m; r++) {          // top-left row
    for (let c = 0; c < n; c++) {        // top-left column
      // greedily grow while a bigger square fits and stays light enough
      while (best + 1 <= Math.min(m - r, n - c) &&
             squareSum(r, c, best + 1) <= threshold) {
        best++;
      }
    }
  }
  return best;
}

int maxSideLength(int[][] mat, int threshold) {
    int m = mat.length, n = mat[0].length;
    // pre[i][j] = sum of mat[0..i-1][0..j-1]
    int[][] pre = new int[m + 1][n + 1];
    for (int i = 1; i <= m; i++)
        for (int j = 1; j <= n; j++)
            pre[i][j] = mat[i - 1][j - 1] + pre[i - 1][j]
                      + pre[i][j - 1] - pre[i - 1][j - 1];

    int best = 0;
    for (int r = 0; r < m; r++) {            // top-left row
        for (int c = 0; c < n; c++) {        // top-left column
            // grow while a bigger square fits and stays light enough
            while (best + 1 <= Math.min(m - r, n - c)
                   && squareSum(pre, r, c, best + 1) <= threshold) {
                best++;
            }
        }
    }
    return best;
}

int squareSum(int[][] pre, int r, int c, int k) {   // side-k square sum
    return pre[r + k][c + k] - pre[r][c + k]
         - pre[r + k][c] + pre[r][c];
}

int squareSum(vector<vector<int>>& pre, int r, int c, int k) {
    return pre[r + k][c + k] - pre[r][c + k]
         - pre[r + k][c] + pre[r][c];
}

int maxSideLength(vector<vector<int>>& mat, int threshold) {
    int m = mat.size(), n = mat[0].size();
    // pre[i][j] = sum of mat[0..i-1][0..j-1]
    vector<vector<int>> pre(m + 1, vector<int>(n + 1, 0));
    for (int i = 1; i <= m; i++)
        for (int j = 1; j <= n; j++)
            pre[i][j] = mat[i - 1][j - 1] + pre[i - 1][j]
                      + pre[i][j - 1] - pre[i - 1][j - 1];

    int best = 0;
    for (int r = 0; r < m; r++)              // top-left row
        for (int c = 0; c < n; c++)          // top-left column
            // grow while a bigger square fits and stays light enough
            while (best + 1 <= min(m - r, n - c)
                   && squareSum(pre, r, c, best + 1) <= threshold)
                best++;
    return best;
}

Explanation

Naively summing every possible square would be far too slow, so we precompute a 2D prefix-sum table pre with one extra row and column of zeros. pre[i][j] stores the sum of the entire block of original cells from (0, 0) up to (i-1, j-1).

Once that table is built, the sum of any axis-aligned rectangle is a constant-time inclusion–exclusion: bottom-right minus the strip above minus the strip to the left, plus the corner we subtracted twice. For a square with top-left (r, c) and side k that is exactly pre[r+k][c+k] - pre[r][c+k] - pre[r+k][c] + pre[r][c].

To find the largest valid side we scan every cell as a candidate top-left. We keep a running best answer and, at each cell, greedily try to grow: while a square of side best + 1 still fits inside the matrix and its sum is within threshold, we bump best. Because best only ever increases, the total growth work is bounded and the scan stays efficient.

This monotonicity (if side k works somewhere, smaller sides work too) is exactly what makes binary search on the side length an alternative; the greedy-grow scan achieves the same answer in a single pass.

Example: mat with three identical rows and threshold = 4. The lightest 2×2 block in the top-left sums to 4, which is allowed, but no 3×3 square stays within 4, so the answer is 2.

Time: O(m · n) Space: O(m · n)