Maximum number of k divisible components

Tree Depth-First Search

74.1% acceptance

Feb 25, 2026

946

There is an undirected tree with n nodes labeled from 0 to n - 1. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree. You are also given a 0-indexed integer array values of length n, where values[i] is the value associated with the ith node, and an integer k. A valid split of the tree is obtained by removing any set of edges, possibly empty, from the tree such that the resulting components all have values that are divisible by k, where the value of a connected component is the sum of the values of its nodes. Return the maximum number of components in any valid split.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

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 k = k as i64;
    let mut adj = vec![vec![]; n];
    for e in &edges {
      let (a, b) = (e[0] as usize, e[1] as usize);
      adj[a].push(b);
      adj[b].push(a);
    }
    // DFS: compute subtree sums. If subtree sum % k == 0, cut edge above (count component)
    let mut vals: Vec<i64> = values.iter().map(|&v| v as i64).collect();
    let mut ans = 0i32;
    let mut parent = vec![n; n];
    let mut order = vec![];
    let mut stack = vec![(0usize, n)];
    while let Some((u, p)) = stack.pop() {
      order.push((u, p));
      parent[u] = p;
      for &v in &adj[u] {
        if v != p { stack.push((v, u)); }
      }
    }
    for &(u, p) in order.iter().rev() {
      if vals[u] % k == 0 {
        ans += 1;
        vals[u] = 0;
      }
      if p != n {
        vals[p] = (vals[p] + vals[u]) % k;
      }
    }
    ans
  }
}