1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
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,
)
)
|