LeetCode Solutions Blog

Count Subtrees With Max Distance Between Cities

Number: 1740

Difficulty: Hard

Paid? No

Companies: N/A

Problem Description

Given n cities (numbered from 1 to n) connected by n-1 edges forming a tree, count, for every d from 1 to n-1, the number of subtrees (i.e. connected subsets of nodes) whose maximum distance (measured as number of edges in the unique path between two cities) is exactly d.

Key Insights

Since n is at most 15, it is feasible to enumerate every subset (using bitmasking) in O(2^n) time.
For each subset, first check connectivity (using DFS or BFS over nodes in the subset).
For connected subsets with at least two nodes, compute the tree’s diameter (the maximum distance between any two nodes). This can be done via:
- Running two BFS/DFS traversals (first to find an extreme node, then to measure the distance from that extreme).
- Or performing a pairwise check (acceptable since n is small).
Increment the count for the corresponding maximum distance if it falls in the range [1, n-1].

Space and Time Complexity

Time Complexity: O(2^n * (n + n)), roughly O(2^n * n) where n is up to 15. Space Complexity: O(n + 2^n) due to the recursion/queue space and storing the graph (though the dominant factor is the enumeration loop over bitmasks).

Solution

We solve the problem by enumerating each possible subset of nodes using bitmasking. For each subset that contains at least two cities, we perform the following steps:

Check connectivity by selecting an arbitrary node in the subset and using DFS to see if all nodes in the subset can be reached.
If the subset is connected, compute the diameter of the subtree. One efficient approach is:
- Pick any node in the subset and perform BFS to find the farthest node from it.
- Then perform BFS starting from that farthest node to compute the maximum distance (the diameter).
Use the computed diameter to update the answer count for that specific distance. This use of bitmask enumeration and exploiting tree properties allows solving the problem within the small constraint of n ≤ 15.

Code Solutions

# Python solution with line-by-line comments

from collections import deque

def countSubtreesWithMaxDistance(n, edges):
    # Build the graph as an adjacency list (using 0-indexed nodes)
    graph = [[] for _ in range(n)]
    for u, v in edges:
        # Convert to 0-indexed
        graph[u-1].append(v-1)
        graph[v-1].append(u-1)

    # Function to check if a subset (mask) is connected
    def is_connected(mask):
        # Find a starting node (any node in the subset)
        start = None
        for i in range(n):
            if mask & (1 << i):
                start = i
                break
        # BFS from the starting node
        q = deque([start])
        seen = 1 << start
        while q:
            u = q.popleft()
            for v in graph[u]:
                # Check if v is in the subset and not visited
                if (mask & (1 << v)) and not (seen & (1 << v)):
                    seen |= (1 << v)
                    q.append(v)
        # Subset is connected if all nodes in mask are visited
        return seen == mask

    # BFS routine to compute distances from a starting node,
    # only considering nodes inside the given "mask" (subset)
    def bfs(start, mask):
        dist = [-1] * n
        q = deque([start])
        dist[start] = 0
        while q:
            u = q.popleft()
            for v in graph[u]:
                if mask & (1 << v) and dist[v] == -1:
                    dist[v] = dist[u] + 1
                    q.append(v)
        return dist

    # Initialize the answer array for distances d=1 to n-1.
    answer = [0] * (n - 1)
    
    # Iterate over every subset (bitmask) of vertices from 0 to 2^n - 1.
    # Only consider subtrees with at least 2 nodes.
    for mask in range(1, 1 << n):
        # Count nodes in this subset using bit count.
        if bin(mask).count("1") < 2:
            continue
        # Check if the subset is connected.
        if not is_connected(mask):
            continue
        
        # Find an arbitrary node in the subset.
        start_node = 0
        while not (mask & (1 << start_node)):
            start_node += 1

        # First BFS: from start_node, find the farthest node.
        dist_from_start = bfs(start_node, mask)
        farthest = start_node
        for i in range(n):
            if mask & (1 << i) and dist_from_start[i] > dist_from_start[farthest]:
                farthest = i

        # Second BFS: from the farthest node, determine the diameter.
        dist_from_far = bfs(farthest, mask)
        diameter = max(dist_from_far[i] for i in range(n) if mask & (1 << i))
        
        # If diameter is within valid range (>=1), update the answer.
        if 1 <= diameter <= n - 1:
            answer[diameter - 1] += 1

    return answer

# Example usage:
print(countSubtreesWithMaxDistance(4, [[1,2],[2,3],[2,4]]))  # Expected output: [3,4,0]