2) See what the maximum distance is for each subtree.
To do 1), we use a bitmap to represent the graph. The nth bit is 1 if city n is in the graph, and 0 otherwise. This way we iterate through all the possible subtrees by going from 1 to 2^n.
To do 2), we do a bfs from every city in the subtree and use that to find the farthest away node.
Once we do have the maximum distance for a subtree, we can increment the counts for all d less than or equal to that in the answer array.
class Solution(object):
def countSubgraphsForEachDiameter(self, n, edges):
adj_list = defaultdict(list)
for u, v in edges:
adj_list[u-1].append(v-1)
adj_list[v-1].append(u-1)
def bfs(src, cities):
visited = set([src])
q = deque([(src, 0)]) # Pair of (vertex, distance)
farthestDist = 0 # Farthest distance from src to other nodes
while len(q) > 0:
u, d = q.popleft()
farthestDist = d
for v in adj_list[u]:
if v not in visited and v in cities:
visited.add(v)
q.append((v, d+1))
return farthestDist, visited
def maxDistance(tree_bitmask): # return: maximum distance between any two cities in our subset. O(n^2)
cities = set()
for i in range(n):
# Check if the ith bit is set
if (tree_bitmask >> i) & 1 == 1:
cities.add(i)
ans = 0
for i in cities:
farthestDist, visited = bfs(i, cities)
if len(visited) < len(cities): return 0 # Can't visit all nodes of the tree -> Invalid tree
ans = max(ans, farthestDist)
return ans
ans = [0] * (n - 1)
for subtree_candidate in range(1, 2 ** n):
d = maxDistance(subtree_candidate)
if d > 0: ans[d - 1] += 1
return ans