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

Question

CleanShot 2024-10-07 at 10.19.46.png

Code

class Solution:
    def maxSideLength(self, mat: List[List[int]], threshold: int) -> int:
        ROW, COL = len(mat), len(mat[0])

        

        aggregate = [[0] * COL for _ in range(ROW)]

        aggregate[0][0] = mat[0][0]

        for r in range(1, ROW):
            aggregate[r][0] = aggregate[r - 1][0] + mat[r][0]

        for c in range(1, COL):
            aggregate[0][c] = aggregate[0][c - 1] + mat[0][c]


        for r in range(1, ROW):
            for c in range(1, COL):
                aggregate[r][c] = aggregate[r - 1][c] + aggregate[r][c - 1] - aggregate[r - 1][c - 1] + mat[r][c]
        
        def can_form_square(m):

            if m == 0:
                return True


            for r in range(m - 1, ROW):
                for c in range(m - 1, COL):
                    prev_row, prev_col = 0, 0
                    prev_inter = 0

                    if r - m >= 0:
                        prev_row = aggregate[r - m][c]

                    if c - m >= 0:
                        prev_col = aggregate[r][c - m]

                    if r - m >= 0 and c - m >= 0:
                        prev_inter = aggregate[r - m][c - m]

                    cur_square = aggregate[r][c] - prev_row - prev_col + prev_inter
                    if cur_square <= threshold:
                        return True
            return False


        l, r = 0, min(ROW, COL) + 1

        while l < r:
            m = (l + r) // 2

            if can_form_square(m):
                l = m + 1
            else:
                r = m

        return l - 1