Count Subtrees With Max Distance Between Cities LeetCode solution in Python

1
# Runtime: 2577 ms (Top 35.48%) | Memory: 14.1 MB (Top 20.97%)
2
from collections import defaultdict
3

4

5
class Solution:
6
    def countSubgraphsForEachDiameter(
7
        self, n: int, edges: List[List[int]]
8
    ) -> List[int]:
9
        def floydwarshall(n, edges):
10
            D = [[float("inf")] * n for _ in range(n)]
11
            for u, v in edges:
12
                D[u - 1][v - 1] = 1
13
                D[v - 1][u - 1] = 1
14
            for v in range(n):
15
                D[v][v] = 0
16
            for k in range(n):
17
                for i in range(n):
18
                    for j in range(n):
19
                        D[i][j] = min(D[i][j], D[i][k] + D[k][j])
20
            return D
21

22
        G = defaultdict(list)
23
        for v, w in edges:
24
            G[v - 1].append(w - 1)
25
            G[w - 1].append(v - 1)
26

27
        def dfs(G, v, V, visited):
28
            if v in visited:
29
                return
30
            visited.add(v)
31
            for w in G[v]:
32
                if w in V:
33
                    dfs(G, w, V, visited)
34

35
        def check_connected(G, V):
36
            if not V:
37
                return False
38
            root = next(iter(V))
39
            visited = set()
40
            dfs(G, root, V, visited)
41
            return visited == V
42

43
        def max_distance(D, V):
44
            res = 0
45
            for v in V:
46
                for w in V:
47
                    res = max(res, D[v][w])
48
            return res
49

50
        def solve(include, idx, ans, G, D):
51
            if idx == n:
52
                V = set(v for v in range(n) if include[v])
53
                if check_connected(G, V):
54
                    d = max_distance(D, V)
55
                    if d >= 1:
56
                        ans[d - 1] += 1
57
            else:
58
                solve(include, idx + 1, ans, G, D)
59
                include[idx] = True
60
                solve(include, idx + 1, ans, G, D)
61
                include[idx] = False
62

63
        D = floydwarshall(n, edges)
64
        include = [False] * n
65
        ans = [0] * (n - 1)
66
        solve(include, 0, ans, G, D)
67
        return ans