Find sorted submatrices with maximum element at most k

Array Stack Matrix Monotonic Stack

50.8% acceptance

Mar 31, 2026

You are given a 2D matrix grid of size m x n. You are also given a non-negative integer k. Return the number of submatrices of grid that satisfy the following conditions: The maximum element in the submatrix less than or equal to k. Each row in the submatrix is sorted in non-increasing order. A submatrix (x1, y1, x2, y2) is a matrix that forms by choosing all cells grid[x][y] where x1 <= x <= x2 and y1 <= y <= y2.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn count_submatrices(grid: Vec<Vec<i32>>, k: i32) -> i64 {
    let m = grid.len();
    let n = grid[0].len();
    // For each cell (i,j), compute w[i][j] = max width of non-increasing row ending at (i,j) with all values <= k
    let mut w = vec![vec![0i32; n]; m];
    for i in 0..m {
      for j in 0..n {
        if grid[i][j] > k {
          w[i][j] = 0;
        } else if j == 0 || grid[i][j] > grid[i][j - 1] {
          w[i][j] = 1;
        } else {
          w[i][j] = w[i][j - 1] + 1;
        }
      }
    }
    // For each column j, go down rows. For a submatrix ending at (i,j) with height h,
    // the width is min(w[i-h+1..=i][j]).
    // Use a stack-based approach: for each (i,j), the valid submatrix count = 
    // sum over h of min(w[i-h+1..=i][j])
    // We can compute this column by column using a monotone stack.
    let mut ans: i64 = 0;
    for j in 0..n {
      // Stack stores (width, count): repeated entries of same width
      let mut stack: Vec<(i64, i64)> = Vec::new();
      let mut total: i64 = 0; // sum of all width*count in the stack
      for i in 0..m {
        let wij = w[i][j] as i64;
        if wij == 0 {
          stack.clear();
          total = 0;
        } else {
          let mut cnt: i64 = 1;
          // Pop elements from stack that have width >= wij
          while let Some(&(top_w, top_c)) = stack.last() {
            if top_w >= wij {
              stack.pop();
              total -= top_w * top_c;
              cnt += top_c;
            } else {
              break;
            }
          }
          stack.push((wij, cnt));
          total += wij * cnt;
          ans += total;
        }
      }
    }
    ans
  }
}