Hard
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.
Example 1:
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:
It can be shown that no other valid split has more than 2 connected components.
Example 2:
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:
It can be shown that no other valid split has more than 3 connected components.
Constraints:
1 <= n <= 3 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
values.length == n
0 <= values[i] <= 109
1 <= k <= 109
values
is divisible by k
.edges
represents a valid tree.class Solution {
private var ans = 0
fun maxKDivisibleComponents(n: Int, edges: Array<IntArray>, values: IntArray, k: Int): Int {
val adj: MutableList<MutableList<Int>> = ArrayList()
for (i in 0 until n) {
adj.add(ArrayList())
}
for (edge in edges) {
val start = edge[0]
val end = edge[1]
adj[start].add(end)
adj[end].add(start)
}
val isVis = BooleanArray(n)
isVis[0] = true
get(0, -1, adj, isVis, values, k.toLong())
return ans
}
private fun get(
curNode: Int,
parent: Int,
adj: List<MutableList<Int>>,
isVis: BooleanArray,
values: IntArray,
k: Long,
): Long {
var sum = values[curNode].toLong()
for (ele in adj[curNode]) {
if (ele != parent && !isVis[ele]) {
isVis[ele] = true
sum += get(ele, curNode, adj, isVis, values, k)
}
}
return if (sum % k == 0L) {
ans++
0
} else {
sum
}
}
}