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

Problem #

Given a m x n matrix mat and an integer threshold, return _the maximum side-length of a square with a sum less than or equal to_threshold _or return_0 if there is no such square.

Example 1:

Input: mat = [[1,1,3,2,4,3,2],[1,1,3,2,4,3,2],[1,1,3,2,4,3,2]], threshold = 4
Output: 2
Explanation: The maximum side length of square with sum less than or equal to 4 is 2 as shown.

Example 2:

Input: mat = [[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2]], threshold = 1
Output: 0

Constraints:

m == mat.length
n == mat[i].length
1 <= m, n <= 300
0 <= mat[i][j] <= 10^4
0 <= threshold <= 10^5

Code #

package leetcode

// Solution 1: 2D prefix sum. sum[i+1][j+1] represents the sum of the rectangle from the top-left corner to (i,j).
// The side length ans starts from 0 and increases: each time, check whether there exists a square with side length ans+1
// and sum not exceeding threshold; if it exists, increment ans by one. Since the answer is monotonic, this is equivalent
// to a single scan overall, with time complexity O(m*n).
func maxSideLength(mat [][]int, threshold int) int {
	m, n := len(mat), len(mat[0])
	sum := make([][]int, m+1)
	for i := range sum {
		sum[i] = make([]int, n+1)
	}
	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			sum[i+1][j+1] = sum[i][j+1] + sum[i+1][j] - sum[i][j] + mat[i][j]
		}
	}
	ans := 0
	for i := 1; i <= m; i++ {
		for j := 1; j <= n; j++ {
			// Try a square with (i,j) as the bottom-right corner and side length ans+1
			k := ans + 1
			if i >= k && j >= k {
				area := sum[i][j] - sum[i-k][j] - sum[i][j-k] + sum[i-k][j-k]
				if area <= threshold {
					ans++
				}
			}
		}
	}
	return ans
}