3070. 元素和小于等于 k 的子矩阵的数目

src/main.rs

@@ -1,35 +1,39 @@
-use std::cmp::max;
-
 struct Solution;
 mod arr;
 impl Solution {
-    pub fn largest_submatrix(mut matrix: Vec<Vec<i32>>) -> i32 {
-        for i in 1..matrix.len() {
-            for j in 0..matrix[i].len() {
-                if matrix[i][j] == 1 {
-                    matrix[i][j] = 1 + matrix[i - 1][j];
-                }
+    pub fn count_submatrices(mut grid: Vec<Vec<i32>>, k: i32) -> i32 {
+        let mut ans = 0;
+        if grid[0][0] <= k {
+            ans += 1;
+        }else {
+            return 0;
+        }
+        for i in 1..grid[0].len() {
+            grid[0][i] += grid[0][i-1];
+            if grid[0][i] <= k {
+                ans += 1;
             }
         }
-        let mut ans = 0;
-        let mut arr = Vec::with_capacity(matrix[0].len());
-        for i in 0..matrix.len() {
-            for j in 0..matrix[i].len() {
-                arr.push(matrix[i][j]);
+        for i in 1.. grid.len() {
+            grid[i][0] += grid[i-1][0];
+            if grid[i][0] <= k {
+                ans += 1;
             }
-            arr.sort_unstable_by(|a, b| b.cmp(a));
-            for j in 0..matrix[i].len() {
-                ans = max(ans, arr[j] * (j + 1) as i32);
+            for j in 1.. grid[i].len() {
+                grid[i][j] += grid[i-1][j] + grid[i][j-1] - grid[i-1][j-1];
+                if grid[i][j] <= k {
+                    ans += 1;
+                }
             }
-            arr.clear()
         }
         ans
     }
 }
 
 fn main() {
-    let result = Solution::largest_submatrix(arr::make_matrix(
-        "[[0,0,1],[1,1,1],[1,0,1]]",
-    ));
+    let result = Solution::count_submatrices(arr::make_matrix(
+        "[[7,2,9],[1,5,0],[2,6,6]]"
+    ), 20);
+
     println!("{:?}", result);
 }