Maximum Number of K-Divisible Components

hard tree DFS subtree sum

Problem

An undirected tree has n nodes (labels 0..n−1) with the given edges, a node-value array values, and an integer k. A valid split removes any set of edges so every resulting connected component has a node-value sum divisible by k. Return the maximum number of components over all valid splits. The total of values is guaranteed divisible by k.

Inputn = 5, edges = [[0,2],[1,2],[1,3],[2,4]], values = [1,8,1,4,4], k = 6

Output2

Cut the edge 1–2. Component {1,3} sums to 8+4 = 12 (÷6), component {0,2,4} sums to 1+1+4 = 6 (÷6). No split yields more than 2 components.

edges (a:b, comma-separated)

values (comma-separated)

def maxKDivisibleComponents(n, edges, values, k):
    g = [[] for _ in range(n)]            # adjacency list
    for a, b in edges:
        g[a].append(b)
        g[b].append(a)
    count = 0

    def dfs(u, parent):
        nonlocal count
        s = values[u] % k                 # this node's residue
        for v in g[u]:
            if v != parent:
                s += dfs(v, u)            # add child subtree residue
        if s % k == 0:
            count += 1                    # cut off: own component
            return 0                      # contributes 0 to parent
        return s % k                      # carry residue up to parent

    dfs(0, -1)                            # root the tree at node 0
    return count

function maxKDivisibleComponents(n, edges, values, k) {
  const g = Array.from({ length: n }, () => []); // adjacency list
  for (const [a, b] of edges) { g[a].push(b); g[b].push(a); }
  let count = 0;

  function dfs(u, parent) {
    let s = values[u] % k;                // this node's residue
    for (const v of g[u]) {
      if (v !== parent) s += dfs(v, u);   // add child subtree residue
    }
    if (s % k === 0) {
      count++;                            // cut off: own component
      return 0;                           // contributes 0 to parent
    }
    return s % k;                         // carry residue up to parent
  }

  dfs(0, -1);                             // root the tree at node 0
  return count;
}

int count = 0;
public int maxKDivisibleComponents(int n, int[][] edges, int[] values, int k) {
    List<List<Integer>> g = new ArrayList<>();
    for (int i = 0; i < n; i++) g.add(new ArrayList<>());
    for (int[] e : edges) { g.get(e[0]).add(e[1]); g.get(e[1]).add(e[0]); }
    dfs(0, -1, g, values, k);             // root the tree at node 0
    return count;
}

private long dfs(int u, int parent, List<List<Integer>> g, int[] values, int k) {
    long s = values[u] % k;               // this node's residue
    for (int v : g.get(u)) {
        if (v != parent) s += dfs(v, u, g, values, k); // child residue
    }
    if (s % k == 0) {
        count++;                          // cut off: own component
        return 0;                         // contributes 0 to parent
    }
    return s % k;                         // carry residue up to parent
}

int maxKDivisibleComponents(int n, vector<vector<int>>& edges, vector<int>& values, int k) {
    vector<vector<int>> g(n);              // adjacency list
    for (auto& e : edges) { g[e[0]].push_back(e[1]); g[e[1]].push_back(e[0]); }
    int count = 0;

    function<long long(int,int)> dfs = [&](int u, int parent) -> long long {
        long long s = values[u] % k;      // this node's residue
        for (int v : g[u])
            if (v != parent) s += dfs(v, u); // add child subtree residue
        if (s % k == 0) {
            count++;                      // cut off: own component
            return 0LL;                   // contributes 0 to parent
        }
        return s % k;                     // carry residue up to parent
    };

    dfs(0, -1);                           // root the tree at node 0
    return count;
}

Explanation

We root the tree at node 0 and process it with a post-order DFS, deciding for every node whether the edge to its parent should be cut.

For a node u, look at its subtree residue: the sum of all values in u's subtree, taken modulo k. If that residue is 0, the whole subtree already forms a valid component on its own — so we cut the edge to the parent, increment count, and report 0 upward (this subtree contributes nothing more to the parent's running sum). If the residue is nonzero, the subtree cannot stand alone, so we must keep it joined to the parent and pass the residue up.

Because we only need sums modulo k, every returned value stays in 0 … k−1, which keeps the arithmetic small even though raw values can reach 10⁹.

Why is this greedy optimal? A subtree can be detached exactly when its sum is divisible by k. Detaching as early as possible (deepest subtree first) never blocks a later cut, because the residue handed to the parent is the same whether or not a divisible descendant subtree was already split off. So cutting every divisible subtree we encounter maximizes the component count. The guarantee that the total is divisible by k ensures the final cut at the root is always valid.

Example: values = [1,8,1,4,4], k = 6. The subtree under node 1 (nodes 1 and 3) sums to 12, residue 0 → cut, count becomes 1. The remaining nodes 0, 2, 4 sum to 6, residue 0 at the root → cut, count becomes 2. Answer: 2.

Time: O(n) Space: O(n)