2872. 可以被 K 整除连通块的最大数目

src/main.rs

@@ -1,30 +1,44 @@
 use crate::arr::make_matrix;
+use std::collections::HashSet;
 
 struct Solution;
 mod arr;
 impl Solution {
-    pub fn number_of_paths(grid: Vec<Vec<i32>>, k: i32) -> i32 {
-        let mut arr = vec![vec![vec![0; k as usize]; grid[0].len()]; grid.len()];
-        let r#mod = 1e9 as i32 + 7;
-        arr[0][0][(grid[0][0] % k) as usize] = 1;
-        for i in 0..grid.len() {
-            for j in 0..grid[i].len() {
-                if i < grid.len() - 1 {
-                    for l in 0..k {
-                        arr[i + 1][j][((l + grid[i + 1][j]) % k) as usize] = (arr[i + 1][j][((l + grid[i + 1][j]) % k) as usize] + arr[i][j][l as usize]) % r#mod;
-                    }
-                }
-                if j < grid[0].len() - 1 {
-                    for l in 0..k {
-                        arr[i][j + 1][((l + grid[i][j + 1]) % k) as usize] = (arr[i][j + 1][((l + grid[i][j + 1]) % k) as usize] + arr[i][j][l as usize]) % r#mod;
-                    }
-                }
+    pub fn max_k_divisible_components(
+        n: i32,
+        edges: Vec<Vec<i32>>,
+        values: Vec<i32>,
+        k: i32,
+    ) -> i32 {
+        let mut edge = vec![HashSet::new(); n as usize];
+        edges.iter().for_each(|x| {
+            edge[x[0] as usize].insert(x[1] as usize);
+            edge[x[1] as usize].insert(x[0] as usize);
+        });
+        fn dfs(total_ans:&mut i32, k:i32, root: usize, last: usize, edge: &Vec<HashSet<usize>>,values: &Vec<i32>) -> i64{
+            let mut ans = values[root] as i64;
+            for &t in edge[root].iter() {
+                if t == last { continue; }
+                ans += dfs(total_ans,k,t, root, edge,values);
+            }
+            if ans % k as i64 == 0 {
+                *total_ans += 1;
+                0
+            }else {
+                ans
             }
         }
-        arr[grid.len() - 1][grid[0].len() - 1][0]
+        let mut ans = 0;
+        dfs(&mut ans,k,0, usize::MAX, &edge,&values);
+        ans
     }
 }
 fn main() {
-    let result = Solution::number_of_paths(make_matrix("[[5,2,4],[3,0,5],[0,7,2]]"), 3);
+    let result = Solution::max_k_divisible_components(
+        7,
+        make_matrix("[[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]"),
+        vec![3, 0, 6, 1, 5, 2, 1],
+        3,
+    );
     println!("{:?}", result);
 }