# 1504. Count Submatrices With All Ones
"""
given an 'm x n' binary matrix 'mat', return the number of submatrices that
have all ones.
"""
class Solution(object):
def numSubmat(self, mat):
"""
:type mat: List[List[int]]
:rtype: int
"""
n, m = len(mat), len(mat[0])
dp, ans = [0] * m, 0
for i in range(n):
for j in range(m):
dp[j] = dp[j] + 1 if mat[i][j] else 0
total, st = [0] * m, []
for j in range(m):
while st and dp[st[-1]] >= dp[j]:
st.pop()
if st:
curr = st[-1]
total[j] = total[curr] + dp[j] * (j - curr)
else:
total[j] = dp[j] * (j + 1)
st.append(j)
ans += total[j]
return ans
if __name__ == "__main__":
obj = Solution()
print(obj.numSubmat(mat=[[1, 0, 1], [1, 1, 0], [1, 1, 0]]))
print(obj.numSubmat(mat=[[0, 1, 1, 0], [0, 1, 1, 1], [1, 1, 1, 0]]))
