Bit manipulation and diameter of a tree

1617. Count Subtrees With Max Distance Between Cities

Problem

Explanation

bit manipulation to traverse all the subsets

Check whether a subset is a Tree

If it is a tree, get the diameter.

Summary

Code

class Solution:
# return whehter the graph is a tree, the last node visited by BFS
# and the max depth of the tree
def bfs(self, node, graph):
Q = collections.deque([(node, 0)])
last_node = None
max_depth = 0
seen = {node}
while Q:
node, depth = Q.popleft()
last_node = node
max_depth = depth
for neighbor in graph[node]:
if neighbor not in seen:
seen.add(neighbor)
Q.append((neighbor, depth + 1))
is_tree = len(seen) == len(graph)
return is_tree, last_node, max_depth

def find_diameter(self, graph):
nodes = list(graph.keys())
is_tree, last_node , _ = self.bfs(nodes[0], graph)
if not is_tree:return -1
_, _, max_depth = self.bfs(last_node, graph)
return max_depth

def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
graph = collections.defaultdict(set)
for a, b in edges:
graph[a].add(b)
graph[b].add(a)
res = collections.defaultdict(int)
for k in range(1, 1<<n):
nodes = set()
for node in range(n):
if (1<<node) & k:
nodes.add(node+1)
if len(nodes)==1:continue
new_graph = collections.defaultdict(set)
for node in nodes:
new_graph[node] = graph[node] & nodes
if not new_graph[node]:break
else:
diameter = self.find_diameter(new_graph)
res[diameter] += 1
return [res[diameter] for diameter in range(1,n)]

Data Scientist/MLE/SWE @takemobi