Leetcode 2872. Maximum Number of K-Divisible Components

문제

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.

• 0 부터 n-1 까지 이름붙혀진 노드로 이루어진 무방향 트리가 있다.
• values에는 각 노드들의 weight값이 주어진다.
• 이 트리를 여러 작은 트리로 나눌때 모든 weight의 합이 k에 의해 나누어 질 경우 가장 많이 나눈 트리의 갯수를 리턴하시오

예시

values = [1,8,1,4,4]

• 2개로 나눌수 있다.

제한

• $1 <= n <= 3 * 10^4$
• $edges.length == n - 1$
• $edges[i].length == 2$
• $0 <= a_i, b_i < n$
• $values.length == n$
• $0 <= values[i] <= 10^9$
• $1 <= k <= 10^9$
• 모든 노드의 weight의 합은 k로 나누어진다.
• 반드시 올바른 값의 쌍이 존재한다.

풀이

• 이 문제는 생각보다 간단한데 특정 Node로 부터 DFS를 실행하여 자신을 포함하는 모든 k로 나누어지는 서브트리들의 갯수를 세는것이 곧 정답이 된다.
• 왜냐?

2 - 6
  ㄴ 1

• 일때 2로부터 DFS를 돌려서 k가 3일때 가능한것은 2 + 6 + 1 과 2 + 1 두개이다.
• 이를 통해 해결할수 있다.

class Solution:
    def maxKDivisibleComponents(self, n: int, edges: List[List[int]], values: List[int], k: int) -> int:
        adj = defaultdict(list)

        for node1, node2 in edges:
            adj[node1].append(node2)
            adj[node2].append(node1)

        res = 0 

        def dfs(current: int, parent: int):
            nonlocal res

            total = values[current]

            for child in adj[current]:
                if child == parent:
                    continue
                total += dfs(child, current)

            if total % k == 0:
                res += 1

            return total


        dfs(0, -1)

        return res