3548. 等和矩阵分割 II

src/main.rs

@@ -1,45 +1,175 @@
+use std::collections::HashSet;
+
 struct Solution;
 mod arr;
 impl Solution {
+    fn only_line(grid: &Vec<Vec<i32>>) -> bool {
+        if grid[0].len() == 2 {
+            return grid[0][0] == grid[0][1];
+        }
+        let mut sum = 0_i64;
+        for i in 0..grid[0].len() {
+            sum += grid[0][i] as i64;
+        }
+        let mut lsum = grid[0][0] as i64;
+        if sum - grid[0][grid[0].len() - 1] as i64 == 2 * grid[0][0] as i64
+            || grid[0][0] as i64 * 2 == sum
+        {
+            return true;
+        }
+        for i in 1..grid[0].len() - 1 {
+            lsum += grid[0][i] as i64;
+            if lsum * 2 == sum
+                || lsum * 2 - grid[0][0] as i64 == sum
+                || lsum * 2 - grid[0][i] as i64 == sum
+                || lsum == (sum - lsum - grid[0][grid[0].len() - 1] as i64)
+            {
+                return true;
+            }
+        }
+        lsum * 2 == sum || (lsum - grid[0][0] as i64) * 2 == sum
+    }
+    fn only_col(grid: &Vec<Vec<i32>>) -> bool {
+        if grid.len() == 2 {
+            return grid[0][0] == grid[1][0];
+        }
+        let mut sum = 0_i64;
+        for i in 0..grid.len() {
+            sum += grid[i][0] as i64;
+        }
+        let mut lsum = grid[0][0] as i64;
+        if sum - grid[grid.len() - 1][0] as i64 == 2 * grid[0][0] as i64
+            || grid[0][0] as i64 * 2 == sum
+        {
+            return true;
+        }
+        for i in 1..grid.len() - 1 {
+            lsum += grid[i][0] as i64;
+            if lsum * 2 == sum
+                || lsum * 2 - grid[0][0] as i64 == sum
+                || lsum * 2 - grid[i][0] as i64 == sum
+                || lsum == (sum - lsum - grid[grid.len() - 1][0] as i64)
+            {
+                return true;
+            }
+        }
+        lsum * 2 == sum || (lsum - grid[0][0] as i64) * 2 == sum
+    }
     pub fn can_partition_grid(grid: Vec<Vec<i32>>) -> bool {
+        if grid.len() == 1 {
+            return Self::only_line(&grid);
+        }
+        if grid[0].len() == 1 {
+            return Self::only_col(&grid);
+        }
         let mut sum = 0_i64;
         let mut line = Vec::<i64>::with_capacity(grid.len());
+        let mut col = vec![0_i64; grid[0].len()];
         for i in 0..grid.len() {
             let mut cnt = 0_i64;
             for j in 0..grid[i].len() {
                 cnt += grid[i][j] as i64;
+                col[j] += grid[i][j] as i64;
             }
             sum += cnt;
             line.push(cnt);
         }
-        let mut lsum = line[0];
+        let mut st = HashSet::<i64>::new();
-        for i in 1..line.len() {
+        let mut l_sum = line[0];
-            if lsum * 2 == sum {
+
+        st.insert(grid[0][0] as i64);
+        st.insert(grid[0][grid[0].len() - 1] as i64);
+
+        if line[0] * 2 == sum || st.contains(&(line[0] * 2 - sum)) {
+            return true;
+        }
+        for i in 1..grid[0].len() - 1 {
+            st.insert(grid[0][i] as i64);
+        }
+        for i in 1..grid.len() - 1 {
+            for j in 0..grid[0].len() {
+                st.insert(grid[i][j] as i64);
+            }
+            l_sum += line[i];
+            if l_sum * 2 == sum || st.contains(&(l_sum * 2 - sum)) {
                 return true;
             }
-            lsum += line[i];
         }
-        let mut column = 0_i64;
+
-        for i in 0..grid[0].len() - 1 {
+        let mut st = HashSet::<i64>::new();
-            let mut cnt = 0_i64;
+        let n = grid.len() - 1;
+        let mut l_sum = line[n];
+
+        st.insert(grid[n][0] as i64);
+        st.insert(grid[n][grid[0].len() - 1] as i64);
+
+        if line[n] * 2 == sum || st.contains(&(line[n] * 2 - sum)) {
+            return true;
+        }
+        for i in 1..grid[0].len() - 1 {
+            st.insert(grid[n][i] as i64);
+        }
+        for i in (1..grid.len() - 1).rev() {
+            for j in 0..grid[0].len() {
+                st.insert(grid[i][j] as i64);
+            }
+            l_sum += line[i];
+            if l_sum * 2 == sum || st.contains(&(l_sum * 2 - sum)) {
+                return true;
+            }
+        }
+
+        let mut st = HashSet::<i64>::new();
+        let mut l_sum = col[0];
+
+        st.insert(grid[0][0] as i64);
+        st.insert(grid[grid.len() - 1][0] as i64);
+
+        if col[0] * 2 == sum || st.contains(&(col[0] * 2 - sum)) {
+            return true;
+        }
+        for i in 1..grid.len() - 1 {
+            st.insert(grid[i][0] as i64);
+        }
+        for i in 1..grid[0].len() - 1 {
             for j in 0..grid.len() {
-                cnt += grid[j][i] as i64;
+                st.insert(grid[j][i] as i64);
             }
-            column += cnt;
+            l_sum += col[i];
-            if column * 2 == sum {
+            if l_sum * 2 == sum || st.contains(&(l_sum * 2 - sum)) {
                 return true;
             }
         }
+
+        let mut st = HashSet::<i64>::new();
+        let n = grid[0].len() - 1;
+        let mut l_sum = col[n];
+
+
+        st.insert(grid[0][n] as i64);
+        st.insert(grid[grid.len() - 1][n] as i64);
+
+        if col[n] * 2 == sum || st.contains(&(col[n] * 2 - sum)) {
+            return true;
+        }
+        for i in 1..grid.len() - 1 {
+            st.insert(grid[i][n] as i64);
+        }
+        for i in (1..grid[0].len() - 1).rev() {
+            for j in 0..grid.len() {
+                st.insert(grid[j][i] as i64);
+            }
+            l_sum += col[i];
+            if l_sum * 2 == sum || st.contains(&(l_sum * 2 - sum)) {
+                return true;
+            }
+        }
+
         false
     }
 }
 
 fn main() {
-    let result = Solution::can_partition_grid(arr::make_matrix(r#"[[42047],[57775],[99822]]"#));
+    let result = Solution::can_partition_grid(arr::make_matrix(r#"[[100000,100000,100000,39719,100000],[50786,100000,100000,100000,90505]]"#));
     println!("{:?}", result);
 }
-/*
-[[1,4,4,0]
-,[-2,0,0,1]
-,[1,-1,1,1]]
- */