Maximum Side Length of a Square with Sum Less than or Equal to Threshold LeetCode solution in Python

1
class Solution:
2
    def maxSideLength(self, mat: List[List[int]], threshold: int) -> int:
3
        # prefix matrix
4
        dp = [[0] * (len(mat[0]) + 1) for _ in range(len(mat) + 1)]
5

6
        for i in range(1, len(mat) + 1):
7
            for j in range(1, len(mat[0]) + 1):
8
                dp[i][j] = (
9
                    dp[i][j - 1] + dp[i - 1][j] - dp[i - 1][j - 1] + mat[i - 1][j - 1]
10
                )
11

12
        # bin search
13
        max_side = 0
14
        for i in range(1, len(mat) + 1):
15
            for j in range(1, len(mat[0]) + 1):
16
                if min(i, j) < max_side:
17
                    continue
18
                left = 0
19
                right = min(i, j)
20
                while left <= right:
21
                    mid = (left + right) // 2
22
                    pref_sum = (
23
                        dp[i][j]
24
                        - dp[i - mid][j]
25
                        - dp[i][j - mid]
26
                        + dp[i - mid][j - mid]
27
                    )
28
                    if pref_sum <= threshold:
29
                        max_side = max(max_side, mid)
30
                        left = mid + 1
31
                    else:
32
                        right = mid - 1
33
        return max_side