Number Of Ways To Reconstruct A Tree LeetCode solution in Python

1
class Solution:
2
    def checkWays(self, P):
3
        g = defaultdict(set)
4
        for u, v in P:
5
            g[u].add(v)
6
            g[v].add(u)
7

8
        def helper(nodes):
9
            d, m = defaultdict(list), len(nodes) - 1
10
            for node in nodes:
11
                d[len(g[node])].append(node)
12

13
            if len(d[m]) == 0:
14
                return 0
15
            root = d[m][0]
16

17
            for node in g[root]:
18
                g[node].remove(root)
19

20
            comps, seen, i = defaultdict(set), set(), 0
21

22
            def dfs(node, i):
23
                comps[i].add(node)
24
                seen.add(node)
25
                for neib in g[node]:
26
                    if neib not in seen:
27
                        dfs(neib, i)
28

29
            for node in nodes:
30
                if node != root and node not in seen:
31
                    dfs(node, i)
32
                    i += 1
33

34
            cands = [helper(comps[i]) for i in comps]
35
            if 0 in cands:
36
                return 0
37
            if 2 in cands:
38
                return 2
39
            if len(d[m]) >= 2:
40
                return 2
41
            return 1
42

43
        return helper(set(g.keys()))