aboutsummaryrefslogtreecommitdiff
path: root/count_submatrices_all_ones.py
blob: e022235ab4c3aec345e35dab222104436610cf35 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# 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]]))