Skip to content

1504. Count Submatrices With All Ones 👍

Approach 1: Naive

  • Time: $O(m^2n)$
  • Space: $O(1)$
 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
class Solution {
 public:
  int numSubmat(vector<vector<int>>& mat) {
    const int m = mat.size();
    const int n = mat[0].size();
    int ans = 0;

    for (int baseRow = 0; baseRow < m; ++baseRow) {
      vector<int> row(n, 1);
      for (int i = baseRow; i < m; ++i) {
        for (int j = 0; j < n; ++j)
          row[j] &= mat[i][j];
        ans += count(row);
      }
    }

    return ans;
  }

 private:
  int count(vector<int>& row) {
    int res = 0;
    int length = 0;
    for (const int num : row) {
      length = num == 0 ? 0 : length + 1;
      res += length;
    }
    return res;
  }
};

 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
class Solution {
  public int numSubmat(int[][] mat) {
    final int m = mat.length;
    final int n = mat[0].length;
    int ans = 0;

    for (int baseRow = 0; baseRow < m; ++baseRow) {
      int[] row = new int[n];
      Arrays.fill(row, 1);
      for (int i = baseRow; i < m; ++i) {
        for (int j = 0; j < n; ++j)
          row[j] &= mat[i][j];
        ans += count(row);
      }
    }

    return ans;
  }

  private int count(int[] row) {
    int res = 0;
    int length = 0;
    for (final int num : row) {
      length = num == 0 ? 0 : length + 1;
      res += length;
    }
    return res;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
  def numSubmat(self, mat: list[list[int]]) -> int:
    m = len(mat)
    n = len(mat[0])
    ans = 0

    for baseRow in range(m):
      row = [1] * n
      for i in range(baseRow, m):
        for j in range(n):
          row[j] &= mat[i][j]
        ans += self._count(row)

    return ans

  def _count(self, row: list[int]) -> int:
    res = 0
    length = 0
    for num in row:
      length = 0 if num == 0 else length + 1
      res += length
    return res

Approach 2: Stack

  • Time: $O(mn)$
  • Space: $O(n)$
 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
class Solution {
 public:
  int numSubmat(vector<vector<int>>& mat) {
    int ans = 0;
    vector<int> hist(mat[0].size());

    for (const vector<int>& row : mat) {
      for (int i = 0; i < row.size(); ++i)
        hist[i] = row[i] == 0 ? 0 : hist[i] + 1;
      ans += count(hist);
    }

    return ans;
  }

 private:
  int count(const vector<int>& hist) {
    // submatrices[i] := the number of submatrices, where the i-th column is the
    // right border
    vector<int> submatrices(hist.size());
    stack<int> stack;

    for (int i = 0; i < hist.size(); ++i) {
      while (!stack.empty() && hist[stack.top()] >= hist[i])
        stack.pop();
      if (!stack.empty()) {
        const int prevIndex = stack.top();
        submatrices[i] = submatrices[prevIndex] + hist[i] * (i - prevIndex);
      } else {
        submatrices[i] = hist[i] * (i + 1);
      }
      stack.push(i);
    }

    return accumulate(submatrices.begin(), submatrices.end(), 0);
  }
};

 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
class Solution {
  public int numSubmat(int[][] mat) {
    int ans = 0;
    int[] hist = new int[mat[0].length];

    for (int[] row : mat) {
      for (int i = 0; i < row.length; ++i)
        hist[i] = row[i] == 0 ? 0 : hist[i] + 1;
      ans += count(hist);
    }

    return ans;
  }

  private int count(int[] hist) {
    // submatrices[i] := the number of submatrices, where the i-th column is the
    // right border
    int[] submatrices = new int[hist.length];
    Deque<Integer> stack = new ArrayDeque<>();

    for (int i = 0; i < hist.length; ++i) {
      while (!stack.isEmpty() && hist[stack.peek()] >= hist[i])
        stack.pop();
      if (!stack.isEmpty()) {
        final int prevIndex = stack.peek();
        submatrices[i] = submatrices[prevIndex] + hist[i] * (i - prevIndex);
      } else {
        submatrices[i] = hist[i] * (i + 1);
      }
      stack.push(i);
    }

    return Arrays.stream(submatrices).sum();
  }
}

 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
class Solution:
  def numSubmat(self, mat: list[list[int]]) -> int:
    ans = 0
    hist = [0] * len(mat[0])

    for row in mat:
      for i, num in enumerate(row):
        hist[i] = 0 if num == 0 else hist[i] + 1
      ans += self._count(hist)

    return ans

  def _count(self, hist: list[int]) -> int:
    # submatrices[i] := the number of submatrices, where the i-th column is the
    # right border
    submatrices = [0] * len(hist)
    stack = []

    for i, h in enumerate(hist):
      while stack and hist[stack[-1]] >= h:
        stack.pop()
      if stack:
        prevIndex = stack[-1]
        submatrices[i] = submatrices[prevIndex] + h * (i - prevIndex)
      else:
        submatrices[i] = h * (i + 1)
      stack.append(i)

    return sum(submatrices)