Medium
There is an undirected tree with n
nodes labeled from 0
to n - 1
, and rooted at node 0
. You are given 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.
A node is good if all the subtrees rooted at its children have the same size.
Return the number of good nodes in the given tree.
A subtree of treeName
is a tree consisting of a node in treeName
and all of its descendants.
Example 1:
Input: edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]
Output: 7
Explanation:
All of the nodes of the given tree are good.
Example 2:
Input: edges = [[0,1],[1,2],[2,3],[3,4],[0,5],[1,6],[2,7],[3,8]]
Output: 6
Explanation:
There are 6 good nodes in the given tree. They are colored in the image above.
Example 3:
Input: edges = [[0,1],[1,2],[1,3],[1,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[9,12],[10,11]]
Output: 12
Explanation:
All nodes except node 9 are good.
Constraints:
2 <= n <= 105
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
edges
represents a valid tree.class Solution {
private var count = 0
fun countGoodNodes(edges: Array<IntArray>): Int {
val n = edges.size + 1
val nodes = arrayOfNulls<TNode>(n)
nodes[0] = TNode()
for (edge in edges) {
val a = edge[0]
val b = edge[1]
if (nodes[b] != null && nodes[a] == null) {
nodes[a] = TNode()
nodes[b]!!.children.add(nodes[a])
} else {
if (nodes[a] == null) {
nodes[a] = TNode()
}
if (nodes[b] == null) {
nodes[b] = TNode()
}
nodes[a]!!.children.add(nodes[b])
}
}
sizeOfTree(nodes[0])
return count
}
private fun sizeOfTree(node: TNode?): Int {
if (node!!.size > 0) {
return node.size
}
val children: List<TNode?> = node.children
if (children.isEmpty()) {
count++
node.size = 1
return 1
}
val size = sizeOfTree(children[0])
var sum = size
var goodNode = true
for (i in 1 until children.size) {
val child = children[i]
if (size != sizeOfTree(child)) {
goodNode = false
}
sum += sizeOfTree(child)
}
if (goodNode) {
count++
}
sum++
node.size = sum
return sum
}
private class TNode {
var size: Int = -1
var children: MutableList<TNode?> = ArrayList()
}
}