# 1277. Count Square Submatrices with All Ones
"""
given an 'm * n' matrix of ones and zeros, return how many square submatrices
have all ones.
"""
class Solution(object):
def countSquares(self, matrix):
"""
:type matrix: List[List[int]]
:rtype: int
"""
for i in range(1, len(matrix)):
for j in range(1, len(matrix[0])):
matrix[i][j] += (
min(matrix[i - 1][j], matrix[i][j - 1], matrix[i - 1][j - 1]) + 1
)
return sum(map(sum, matrix))
if __name__ == "__main__":
obj = Solution()
print(obj.countSquares(matrix=[[0, 1, 1, 1], [1, 1, 1, 1], [0, 1, 1, 1]]))
print(obj.countSquares(matrix=[[1, 0, 1], [1, 1, 0], [1, 1, 0]]))
