Bit manipulation and diameter of a tree

1617. Count Subtrees With Max Distance Between Cities

Problem

1617. Count Subtrees With Max Distance Between Cities

Explanation

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.

Summary

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.

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)]

Thanks for reading.