方法一:DFS
我们注意到,题目保证了整棵树的节点值之和可以被 $k$ 整除,因此,如果我们删除一棵元素和能被 $k$ 整除的子树,那么剩下的每个连通块的节点值之和也一定可以被 $k$ 整除。
因此,我们可以使用深度优先搜索的方法,从根节点开始遍历整棵树,对于每个节点,我们计算其子树中所有节点值之和,如果该和能被 $k$ 整除,那么我们就将答案加一。
###python
class Solution:
def maxKDivisibleComponents(
self, n: int, edges: List[List[int]], values: List[int], k: int
) -> int:
def dfs(i: int, fa: int) -> int:
s = values[i]
for j in g[i]:
if j != fa:
s += dfs(j, i)
nonlocal ans
ans += s % k == 0
return s
g = [[] for _ in range(n)]
for a, b in edges:
g[a].append(b)
g[b].append(a)
ans = 0
dfs(0, -1)
return ans
###java
class Solution {
private int ans;
private List<Integer>[] g;
private int[] values;
private int k;
public int maxKDivisibleComponents(int n, int[][] edges, int[] values, int k) {
g = new List[n];
Arrays.setAll(g, i -> new ArrayList<>());
for (int[] e : edges) {
int a = e[0], b = e[1];
g[a].add(b);
g[b].add(a);
}
this.values = values;
this.k = k;
dfs(0, -1);
return ans;
}
private long dfs(int i, int fa) {
long s = values[i];
for (int j : g[i]) {
if (j != fa) {
s += dfs(j, i);
}
}
ans += s % k == 0 ? 1 : 0;
return s;
}
}
###cpp
class Solution {
public:
int maxKDivisibleComponents(int n, vector<vector<int>>& edges, vector<int>& values, int k) {
int ans = 0;
vector<int> g[n];
for (auto& e : edges) {
int a = e[0], b = e[1];
g[a].push_back(b);
g[b].push_back(a);
}
auto dfs = [&](this auto&& dfs, int i, int fa) -> long long {
long long s = values[i];
for (int j : g[i]) {
if (j != fa) {
s += dfs(j, i);
}
}
ans += s % k == 0;
return s;
};
dfs(0, -1);
return ans;
}
};
###go
func maxKDivisibleComponents(n int, edges [][]int, values []int, k int) (ans int) {
g := make([][]int, n)
for _, e := range edges {
a, b := e[0], e[1]
g[a] = append(g[a], b)
g[b] = append(g[b], a)
}
var dfs func(int, int) int
dfs = func(i, fa int) int {
s := values[i]
for _, j := range g[i] {
if j != fa {
s += dfs(j, i)
}
}
if s%k == 0 {
ans++
}
return s
}
dfs(0, -1)
return
}
###ts
function maxKDivisibleComponents(
n: number,
edges: number[][],
values: number[],
k: number,
): number {
const g: number[][] = Array.from({ length: n }, () => []);
for (const [a, b] of edges) {
g[a].push(b);
g[b].push(a);
}
let ans = 0;
const dfs = (i: number, fa: number): number => {
let s = values[i];
for (const j of g[i]) {
if (j !== fa) {
s += dfs(j, i);
}
}
if (s % k === 0) {
++ans;
}
return s;
};
dfs(0, -1);
return ans;
}
###rust
impl Solution {
pub fn max_k_divisible_components(n: i32, edges: Vec<Vec<i32>>, values: Vec<i32>, k: i32) -> i32 {
let n = n as usize;
let mut g = vec![vec![]; n];
for e in edges {
let a = e[0] as usize;
let b = e[1] as usize;
g[a].push(b);
g[b].push(a);
}
let mut ans = 0;
fn dfs(i: usize, fa: i32, g: &Vec<Vec<usize>>, values: &Vec<i32>, k: i32, ans: &mut i32) -> i64 {
let mut s = values[i] as i64;
for &j in &g[i] {
if j as i32 != fa {
s += dfs(j, i as i32, g, values, k, ans);
}
}
if s % k as i64 == 0 {
*ans += 1;
}
s
}
dfs(0, -1, &g, &values, k, &mut ans);
ans
}
}
时间复杂度 $O(n)$,空间复杂度 $O(n)$。其中 $n$ 是树中的节点数。
有任何问题,欢迎评论区交流,欢迎评论区提供其它解题思路(代码),也可以点个赞支持一下作者哈 😄~