Maximum Number of K-Divisible Components

lcpy gen -n 2872   # by problem number
lcpy gen -s maximum_number_of_k_divisible_components   # by problem name

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.

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.

class Solution:
    # Time: O(n)
    # Space: O(n)
    def max_k_divisible_components(
        self, n: int, edges: list[list[int]], values: list[int], k: int
    ) -> int:
        adj: list[list[int]] = [[] for _ in range(n)]
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)

        # Iterative post-order DFS from node 0 (n up to 3 * 10^4, avoid recursion).
        parent = [-1] * n
        order = [0]
        parent[0] = 0
        for node in order:
            for nxt in adj[node]:
                if parent[nxt] == -1:
                    parent[nxt] = node
                    order.append(nxt)

        subtree = values[:]
        count = 0
        for node in reversed(order):
            if subtree[node] % k == 0:
                count += 1
            else:
                subtree[parent[node]] += subtree[node]
        return count