Count Submatrices With All Ones LeetCode solution in Java

1
class Solution {
2
  private int n;
3
  private int res = 0;
4

5
  public int numSubmat(int[][] mat) {
6
    this.n = mat[0].length;
7

8
    // dp[j] : the height (number of consecutive '1's) of column j
9
    int[] dp = new int[n];
10
    for (int i = 0; i < mat.length; i++) {
11
      // calculating (updating) heights
12
      for (int j = 0; j < n; j++) {
13
        dp[j] = mat[i][j] == 1 ? dp[j] + 1 : 0;
14
      }
15
      enumerateRowByMinHeight(dp);
16
    }
17
    return res;
18
  }
19

20
  public void enumerateRowByMinHeight(int[] dp) {
21
    // monotonic stack storing indices : for index p < q in stack, dp[p] < dp[q]
22
    Deque<Integer> stack = new LinkedList<>();
23
    stack.offerLast(-1);
24

25
    for (int j = 0; j < n; j++) {
26
      while (stack.peekLast() != -1 && dp[stack.peekLast()] >= dp[j]) {
27
        int idx = stack.pollLast();
28
        res += dp[idx] * (idx - stack.peekLast()) * (j - idx);
29
      }
30
      stack.offerLast(j);
31
    }
32
    while (stack.peekLast() != -1) {
33
      int idx = stack.pollLast();
34
      res += dp[idx] * (idx - stack.peekLast()) * (n - idx);
35
    }
36
  }
37
}