Bit manipulation and diameter of a tree

1617. Count Subtrees With Max Distance Between Cities



bit manipulation to traverse all the subsets

0 to 2¹⁵

The trick is named bit manipulation

For each number, if the bit value is 1, it mean the subset contains a node.

For example, A, B, C, D, we will have

0000 →{}

0001 →{D}

0010 →{C}

0011 →{C, D}

1111 →{A,B,C,D}

Check whether a subset is a Tree

We can check whether len(seen)== len(nodes)

If it is a tree, get the diameter.


  • traverse all possible subsets
  • for each subset, check whether it is a tree. We can use DFS or BFS, since the diameter of the tree part is using BFS, I am also using BFS here
  • Using two BFS to get the diameter of a tree if the graph is a tree.


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:
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:
res = collections.defaultdict(int)
for k in range(1, 1<<n):
nodes = set()
for node in range(n):
if (1<<node) & k:
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
diameter = self.find_diameter(new_graph)
res[diameter] += 1
return [res[diameter] for diameter in range(1,n)]

Thanks for reading.

Data Scientist/MLE/SWE @takemobi