aboutsummaryrefslogtreecommitdiff
path: root/count_number_root_nodes.py
blob: ddd498253a31302d692af339c1b32db7a0b2bc31 (plain)
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,
        )
    )