aboutsummaryrefslogtreecommitdiff
path: root/count_number_root_nodes.py
diff options
context:
space:
mode:
Diffstat (limited to 'count_number_root_nodes.py')
-rw-r--r--count_number_root_nodes.py66
1 files changed, 66 insertions, 0 deletions
diff --git a/count_number_root_nodes.py b/count_number_root_nodes.py
new file mode 100644
index 0000000..ddd4982
--- /dev/null
+++ b/count_number_root_nodes.py
@@ -0,0 +1,66 @@
+# 2581. Count Number of Possible Root Nodes
+from collections import defaultdict
+
+
+class Solution(object):
+ def rootCount(self, edges, guesses, k):
+ """
+ :type edges: List[List[int]]
+ :type guesses: List[List[int]]
+ :type k: int
+ :rtype: int
+ """
+ adj = defaultdict(set)
+ for u, v in edges:
+ adj[u].add(v)
+ adj[v].add(u)
+ child = defaultdict(set)
+ for u, v in guesses:
+ child[u].add(v)
+ parents, total_guess = {}, 0
+
+ def fill_parent(node, par):
+ nonlocal total_guess
+ parents[node] = par
+ for i in adj[node]:
+ if i != par:
+ if i in child[node]:
+ total_guess += 1
+ fill_parent(i, node)
+
+ fill_parent(0, -1)
+ res = 0
+
+ def dfs(node, par, correct_guess):
+ nonlocal res
+ cur = correct_guess
+ if node in child[par]:
+ cur -= 1
+ if par in child[node]:
+ cur += 1
+ if cur >= k:
+ res += 1
+ for i in adj[node]:
+ if i != par:
+ dfs(i, node, cur)
+
+ dfs(0, -1, total_guess)
+ return res
+
+
+if __name__ == "__main__":
+ obj = Solution()
+ print(
+ obj.rootCount(
+ edges=[[0, 1], [1, 2], [1, 3], [4, 2]],
+ guesses=[[1, 3], [0, 1], [1, 0], [2, 4]],
+ k=3,
+ )
+ )
+ print(
+ obj.rootCount(
+ edges=[[0, 1], [1, 2], [2, 3], [3, 4]],
+ guesses=[[1, 0], [3, 4], [2, 1], [3, 2]],
+ k=1,
+ )
+ )