# Count Square Submatrices with All Ones
# Given a m * n matrix of ones and zeros, return how many square submatrices have all ones.
#
#
# Example 1:
# Input: matrix =
# [
#     [0, 1, 1, 1],
#     [1, 1, 1, 1],
#     [0, 1, 1, 1]
# ]
# Output: 15
# Explanation:
# There are 10 squares of side 1.
# There are 4 squares of side 2.
# There is 1 square of side 3.
# Total number of squares = 10 + 4 + 1 = 15.
#
# Example 2:
# Input: matrix =
# [
#     [1, 0, 1],
#     [1, 1, 0],
#     [1, 1, 0]
# ]
# Output: 7
# Explanation:
# There are 6 squares of side 1.
# There is 1 square of side 2.
# Total number of squares = 6 + 1 = 7.
#
# Constraints:
# 1 <= arr.length <= 300
# 1 <= arr[0].length <= 300
# 0 <= arr[i][j] <= 1

from typing import List


class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        count = 0

        # count all in first row
        for i in range(len(matrix[0])):
            count += matrix[0][i]

        # count all in first col
        for i in range(1, len(matrix)):
            count += matrix[i][0]

        # Now keep on checking for squares and mark them incrementally
        for i in range(1, len(matrix)):
            for j in range(1, len(matrix[0])):
                if matrix[i][j] == 1:
                    n = min(matrix[i-1][j-1], matrix[i-1][j], matrix[i][j-1])
                    if matrix[i][j] <= n:
                        matrix[i][j] = n + 1
                count += matrix[i][j]

        return count
