Number of Submatrices That Sum to Target LeetCode solution in Python

1
class Solution:
2
    def numSubmatrixSumTarget(self, matrix: List[List[int]], target: int) -> int:
3
        m, n = len(matrix), len(matrix[0])
4
        matrix_sums = [[0 for _ in range(n)] for _ in range(m)]
5

6
        # Calculate all the submatrices sum with the transition formula we found
7
        for row in range(m):
8
            for col in range(n):
9
                # first cell
10
                if row == 0 and col == 0:
11
                    matrix_sums[row][col] = matrix[row][col]
12
                # Rows and columns are like prefix sums, without intersection
13
                elif row == 0:
14
                    matrix_sums[row][col] = matrix[row][col] + matrix_sums[row][col - 1]
15
                elif col == 0:
16
                    matrix_sums[row][col] = matrix[row][col] + matrix_sums[row - 1][col]
17

18
                # current sum is the sum of the matrix above, to the left and subtract the intersection
19
                else:
20
                    matrix_sums[row][col] = (
21
                        matrix[row][col]
22
                        + (matrix_sums[row][col - 1])
23
                        + (matrix_sums[row - 1][col])
24
                        - (matrix_sums[row - 1][col - 1])
25
                    )
26

27
        ans = 0
28
        # Generate all submatrices
29
        for y1 in range(m):
30
            for x1 in range(n):
31
                for y2 in range(y1, m):
32
                    for x2 in range(x1, n):
33
                        # calculate sum in O(1)
34
                        submatrix_total = (
35
                            matrix_sums[y2][x2]
36
                            - (matrix_sums[y2][x1 - 1] if x1 - 1 >= 0 else 0)
37
                            - (matrix_sums[y1 - 1][x2] if y1 - 1 >= 0 else 0)
38
                            + (
39
                                matrix_sums[y1 - 1][x1 - 1]
40
                                if y1 - 1 >= 0 and x1 - 1 >= 0
41
                                else 0
42
                            )
43
                        )
44

45
                        if submatrix_total == target:
46
                            ans += 1
47
        return ans