aboutsummaryrefslogtreecommitdiff
path: root/count_num_good_nodes.py
diff options
context:
space:
mode:
Diffstat (limited to 'count_num_good_nodes.py')
-rw-r--r--count_num_good_nodes.py77
1 files changed, 77 insertions, 0 deletions
diff --git a/count_num_good_nodes.py b/count_num_good_nodes.py
new file mode 100644
index 0000000..fa876f9
--- /dev/null
+++ b/count_num_good_nodes.py
@@ -0,0 +1,77 @@
+# 3249. Count the Number of Good Nodes
+from collections import defaultdict
+
+"""
+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] = [a_i, b_i]' indicates that there is an edge between nodes
+'a_i' and 'b_i' in the tree. a node is good if all the subtrees rooted at is
+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.
+"""
+
+
+class Solution(object):
+ def countGoodNodes(self, edges):
+ """
+ :type edges: List[List[int]]
+ :rtype: int
+ """
+ n = len(edges) + 1
+ graph = defaultdict(list)
+ subtree_size = [0] * n
+
+ def dfs(curr, parent):
+ n = 1
+ for i in graph[curr]:
+ if i != parent:
+ n += dfs(i, curr)
+ subtree_size[curr] = n
+ return n
+
+ for u, v in edges:
+ graph[u].append(v)
+ graph[v].append(u)
+ dfs(0, -1)
+ ans = 0
+ for i in range(n):
+ flag, prev = True, -1
+ for j in graph[i]:
+ if subtree_size[j] < subtree_size[i]:
+ if prev == -1:
+ prev = subtree_size[j]
+ elif prev != subtree_size[j]:
+ flag = False
+ break
+ if flag:
+ ans += 1
+ return ans
+
+
+if __name__ == "__main__":
+ obj = Solution()
+ print(obj.countGoodNodes(edges=[[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [2, 6]]))
+ print(
+ obj.countGoodNodes(
+ edges=[[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [1, 6], [2, 7], [3, 8]]
+ )
+ )
+ print(
+ obj.countGoodNodes(
+ 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],
+ ]
+ )
+ )