Bit manipulation and diameter of a tree

1617. Count Subtrees With Max Distance Between Cities

Jimmy (xiaoke) Shen
2 min readOct 11, 2020



1617. Count Subtrees With Max Distance Between Cities


bit manipulation to traverse all the subsets

It only have 15 nodes, if we use bit manipulation, we can list all the possible subset of the 15 nodes by traverse from

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

For each subset, it may not be a tree. For example, if we have AB connected and CD connected. Then A and C can not generate a Tree. Actually, A and C will be a forest. A way to check whether a subset is a Tree is: traverse the graph from any node, if it is a tree, then, the traverse process will collect all the nodes. If not, the traverse process will be interrupted during the process and some nodes will not be seen or visited.

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

If it is a tree, get the diameter.

This is a standard problem and we can use two BFS to solve the problem. For more details see my previous article.


Main steps:

  • 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.



Jimmy (xiaoke) Shen

Data Scientist/MLE/SWE @takemobi