Problem

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 themaximum number of components in any valid split.

Examples

Example 1:

graph LR

    subgraph Original
        A1[1]
        B1[2]
        C1[3]
        D1[4]
        E1[0]
        
        A1 --> B1 & C1
        B1 --> D1 & E1
    end
    
    subgraph After["After Split"]
        A2[1]
        B2[2]
        C2[3]
        D2[4]
        E2[0]

        A2 ~~~ B2
        A2 --> C2
        B2 --> D2 & E2
    end
    
    Original --> After
    
    style A1 stroke:#f66,stroke-width:2px
    style B1 stroke:#f66,stroke-width:2px
  
Input: n = 5, edges = [[0,2],[1,2],[1,3],[2,4]], values = [1,8,1,4,4], k = 6
Output: 2
Explanation: We remove the edge connecting node 1 with 2. The resulting split is valid because:
- The value of the component containing nodes 1 and 3 is values[1] + values[3] = 12.
- The value of the component containing nodes 0, 2, and 4 is values[0] + values[2] + values[4] = 6.
It can be shown that no other valid split has more than 2 connected components.

Example 2:

graph LR

    subgraph Original
        A1[1]
        B1[2]
        C1[3]
        D1[4]
        E1[5]
        F1[6]
        G1[0]
        
        G1 --> B1 & A1
        B1 --> F1 & E1
        A1 --> D1 & C1
    end
    
    subgraph After["After Split"]
        A2[1]
        B2[2]
        C2[3]
        D2[4]
        E2[5]
        F2[6]
        G2[0]
        
        G2 ~~~ B2 & A2
        B2 --> F2 & E2
        A2 --> D2 & C2
    end
    
    Original --> After
    
    style A1 stroke:#f66,stroke-width:2px
    style B1 stroke:#f66,stroke-width:2px
  
Input: n = 7, edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]], values = [3,0,6,1,5,2,1], k = 3
Output: 3
Explanation: We remove the edge connecting node 0 with 2, and the edge connecting node 0 with 1. The resulting split is valid because:
- The value of the component containing node 0 is values[0] = 3.
- The value of the component containing nodes 2, 5, and 6 is values[2] + values[5] + values[6] = 9.
- The value of the component containing nodes 1, 3, and 4 is values[1] + values[3] + values[4] = 6.
It can be shown that no other valid split has more than 3 connected components.

Constraints:

  • 1 <= n <= 3 * 10^4
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • values.length == n
  • 0 <= values[i] <= 10^9
  • 1 <= k <= 10^9
  • Sum of values is divisible by k.
  • The input is generated such that edges represents a valid tree.

Solution

Method 1 - Using DFS

To solve this problem, let’s leverage the structure of a tree. A tree is composed of nodes and edges, where each edge connects a parent node to a child node. When we choose a node as the root, we can decompose the tree into subtrees based on these parent-child relationships. Since the tree is undirected, any node can be selected as the root without affecting the result.

Using recursion, we can calculate the sum of values for each subtree. Once we have the sum, we need to check if it is divisible by k. If it is, we can detach the subtree at that point, treating it as a valid component.

If the sum is not divisible by k, the remainder (i.e., the non-divisible part when divided by k) must be carried over to the parent node. The parent node will then combine its remainder with those of its children to determine if the total sum is divisible by k. This approach is well-suited to Depth-First Search (DFS) recursion:

  • Start from the leaves (smallest subtrees) and calculate their sums.
  • Propagate the results upwards to parent nodes, aggregating remainders modulo k.
  • Count a subtree as a valid component whenever its sum is divisible by k.

Here is the approach:

  1. Initialize an adjacency list adjList to represent the graph.
  2. Populate adjList using the given edges.
  3. Initialize componentCount to 0 to store the count of k-divisible components.
  4. Call dfs(0, -1, adjList, values, k, componentCount) starting from node 0 with no parent (-1).
  5. Return componentCount as the result.

This is how the dfs will work:

  1. Initialize sum to 0, representing the sum of node values in the current subtree.
  2. For each neighborNode of currentNode:
    • If neighborNode is not equal to parentNode, recursively call dfs for neighborNode with currentNode as its parent.
    • Add the result of the recursive call to sum and take modulo k.
  3. Add the value of currentNode to sum and take modulo k.
  4. If sum is 0, increment componentCount because the current subtree forms a k-divisible component.
  5. Return sum to allow the parent node to incorporate the result.

Code

Java
class Solution {
    public int maxKDivisibleComponents(int n, int[][] edges, int[] values, int k) {
        List<Integer>[] adjList = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            adjList[i] = new ArrayList<>();
        }
        for (int[] edge : edges) {
            int node1 = edge[0];
            int node2 = edge[1];
            adjList[node1].add(node2);
            adjList[node2].add(node1);
        }

        int[] componentCount = new int[1];

        dfs(0, -1, adjList, values, k, componentCount);

        return componentCount[0];
    }

    private int dfs(int currentNode, int parentNode, List<Integer>[] adjList, int[] nodeValues, int k, int[] componentCount) {
        int sum = 0;

        for (int neighborNode : adjList[currentNode]) {
            if (neighborNode != parentNode) {
                sum += dfs(neighborNode, currentNode, adjList, nodeValues, k, componentCount);
                sum %= k;
            }
        }

        sum += nodeValues[currentNode];
        sum %= k;

        if (sum == 0) {
            componentCount[0]++;
        }

        return sum;
    }
}
Python
class Solution:
    def maxKDivisibleComponents(self, n: int, edges: List[List[int]], values: List[int], k: int) -> int:
        def dfs(node: int, parent: int) -> int:
            total_sum = 0
            for neighbor in adj[node]:
                if neighbor != parent:
                    total_sum += dfs(neighbor, node)
                    total_sum %= k
            total_sum += values[node]
            total_sum %= k
            if total_sum == 0:
                nonlocal ans
                ans += 1
                return 0
            return total_sum

        adj = [[] for _ in range(n)]
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)

        ans = 0
        dfs(0, -1)
        return ans - 1

Complexity

  • ⏰ Time complexity: O(n), where n is the number of nodes. Each node and edge is visited once during DFS traversal.
  • 🧺 Space complexity: O(n) for storing the adjacency list and the DFS recursion stack.