Count Fertile Pyramids in a Land LeetCode solution in Python

1
class Solution:
2
    def countPyramids(self, grid):
3
        # dp[i][j] represents the number of layers of the largest pyramid with (i, j) as the vertex.
4
        # Note that the 1-level pyramid is invalid in the problem, so it should be removed when summing.
5
        # Note that if grid[i][j] is 0, dp[i][j] will always be 0.
6
        # The dp recurrence formula is dp[i][j] = min(dp[i + 1][j - 1], dp[i + 1][j + 1]) + 1
7
        m, n, dp, cnt = len(grid), len(grid[0]), copy.deepcopy(grid), 0
8
        # triangle
9
        for i in range(m - 2, -1, -1):
10
            for j in range(1, n - 1):
11
                if dp[i][j] > 0 and dp[i + 1][j] > 0:
12
                    dp[i][j] = min(dp[i + 1][j - 1], dp[i + 1][j + 1]) + 1
13
                    cnt += dp[i][j] - 1
14
        # inverted triangle
15
        dp = grid
16
        for i in range(1, m):
17
            for j in range(1, n - 1):
18
                if dp[i][j] > 0 and dp[i - 1][j] > 0:
19
                    dp[i][j] = min(dp[i - 1][j - 1], dp[i - 1][j + 1]) + 1
20
                    cnt += dp[i][j] - 1
21
        return cnt