Problem Description

You are given n cities numbered from 1 to n that are connected by n-1 bidirectional edges, forming a tree structure. This means there's exactly one path between any two cities.

The input consists of:

• An integer n representing the number of cities
• An array edges of size n-1, where each edges[i] = [ui, vi] represents a bidirectional edge between cities ui and vi

A subtree is defined as a subset of cities where:

• Every city in the subset can reach every other city in the subset
• The path between any two cities in the subset only passes through cities that are also in the subset
• Two subtrees are considered different if one contains a city that the other doesn't

The distance between two cities is the number of edges in the path connecting them.

Your task is to count subtrees based on their maximum distance (diameter). For each possible maximum distance d from 1 to n-1, you need to count how many subtrees have exactly d as their maximum distance between any two cities within that subtree.

The output should be an array of size n-1 where:

• The element at index d-1 (since the array is 0-indexed but we count d from 1) represents the number of subtrees that have maximum distance equal to d

For example, if you have 4 cities connected as a tree, you need to find:

• How many subtrees have maximum distance = 1
• How many subtrees have maximum distance = 2
• How many subtrees have maximum distance = 3

The solution involves enumerating all possible non-empty subsets of cities, checking if they form a valid connected subtree, and if so, calculating their diameter (maximum distance between any two cities in the subtree).

Intuition

Since we need to find all possible subtrees and their diameters, and n is small (at most 15 based on typical constraints), we can use a brute force approach with bit manipulation to enumerate all possible subsets of cities.

The key insight is that we can represent any subset of cities using a bitmask, where bit i is set to 1 if city i is included in the subset. For n cities, we have 2^n possible subsets to check.

For each subset (represented by a bitmask), we need to:

1. Check if it forms a valid connected subtree
2. If valid, find its diameter (maximum distance)

To check connectivity and find the diameter simultaneously, we can use a clever DFS approach:

• Start DFS from any city in the subset
• During DFS, track the farthest city reached and the maximum distance
• Toggle off bits in a copy of the mask as we visit cities
• If all bits are toggled off after DFS, the subset is connected

The diameter of a tree can be found using the two-DFS technique:

1. First DFS: Start from any node to find the farthest node from it
2. Second DFS: Start from that farthest node to find the actual diameter

Why does this work? In a tree, the diameter always connects two "peripheral" nodes. The first DFS finds one end of the diameter, and the second DFS from that end gives us the actual diameter length.

We skip single-node subsets (when only one bit is set) because they have no edges and thus no valid distance. This is checked using mask & (mask - 1) == 0, which is true only for powers of 2 (single bit set).

The solution elegantly combines subset enumeration with tree diameter calculation, leveraging the small constraint on n to make brute force feasible while using efficient bit operations and DFS to validate and measure each potential subtree.

Learn more about Tree, Dynamic Programming and Bitmask patterns.

Solution Approach

The implementation uses bit manipulation to enumerate all possible subsets and DFS to validate connectivity and calculate diameters.

Step 1: Build the adjacency list

g = defaultdict(list)
for u, v in edges:
    u, v = u - 1, v - 1  # Convert to 0-indexed
    g[u].append(v)
    g[v].append(u)

We convert the edge list into an adjacency list representation for efficient graph traversal. Cities are converted to 0-indexed for easier bit manipulation.

Step 2: Enumerate all possible subsets

for mask in range(1, 1 << n):
    if mask & (mask - 1) == 0:
        continue

We iterate through all possible bitmasks from 1 to 2^n - 1. The condition mask & (mask - 1) == 0 checks if only one bit is set (power of 2), which we skip since single nodes don't have any edges.

Step 3: DFS function for connectivity check and distance calculation

def dfs(u: int, d: int = 0):
    nonlocal mx, nxt, msk
    if mx < d:
        mx, nxt = d, u
    msk ^= 1 << u
    for v in g[u]:
        if msk >> v & 1:
            dfs(v, d + 1)

The DFS function:

• Tracks the maximum distance mx and the farthest node nxt
• Toggles off the bit for visited node using XOR: msk ^= 1 << u
• Only visits neighbors that are in the current subset: if msk >> v & 1
• Increments distance d as we traverse deeper

Step 4: Validate connectivity and find diameter

msk, mx = mask, 0
cur = msk.bit_length() - 1  # Get the highest set bit (any node in subset)
dfs(cur)
if msk == 0:  # All bits toggled off means all nodes visited (connected)
    msk, mx = mask, 0
    dfs(nxt)  # Second DFS from the farthest node
    ans[mx - 1] += 1

For each subset:

1. Start DFS from any node in the subset (we use the highest bit position)
2. After first DFS, if msk == 0, all nodes were visited (subset is connected)
3. If connected, run second DFS from the farthest node found (nxt) to get the actual diameter
4. Increment the count for the diameter value (mx - 1 because array is 0-indexed)

Key Optimizations:

• Using XOR to toggle bits allows us to check connectivity in one pass
• The two-DFS approach efficiently finds the tree diameter
• Bit manipulation provides compact subset representation and fast operations

The algorithm has time complexity O(2^n * n) where we check each of the 2^n subsets and perform DFS taking O(n) time each.

Example Walkthrough

Let's walk through a small example with n = 4 cities and edges [[1,2], [2,3], [2,4]].

The tree structure looks like:

    1
    |
    2
   / \
  3   4

Step 1: Build adjacency list (0-indexed)

• City 0: [1]
• City 1: [0, 2, 3]
• City 2: [1]
• City 3: [1]

Step 2: Enumerate subsets using bitmasks

Let's trace through a few interesting subsets:

Subset {1, 2} - Bitmask: 0011 (binary)

• Start DFS from city 1
• Visit city 1: toggle bit 1, msk = 0001
• Visit neighbor city 0: toggle bit 0, msk = 0000
• All bits cleared → connected subtree!
• Second DFS from city 0 (farthest found)
• Maximum distance = 1
• Add 1 to ans[0]

Subset {2, 3, 4} - Bitmask: 1110 (binary)

• Start DFS from city 3
• Visit city 3: toggle bit 3, msk = 0110
• Visit neighbor city 1: toggle bit 1, msk = 0100
• Visit neighbor city 2: toggle bit 2, msk = 0000
• All bits cleared → connected subtree!
• Second DFS from city 2 (farthest found)
• Maximum distance = 2 (path from city 2 to city 3/4)
• Add 1 to ans[1]

Subset {1, 3} - Bitmask: 0101 (binary)

• Start DFS from city 2
• Visit city 2: toggle bit 2, msk = 0101 (bit 2 not set initially, no change)
• Can't visit city 1 (not connected without city 2)
• msk ≠ 0 → not a connected subtree, skip!

Subset {1, 2, 3} - Bitmask: 0111 (binary)

• Start DFS from city 2
• Visit city 2: toggle bit 2, msk = 0011
• Visit neighbor city 1: toggle bit 1, msk = 0001
• Visit neighbor city 0: toggle bit 0, msk = 0000
• All bits cleared → connected subtree!
• Second DFS finds maximum distance = 2 (path from city 0 to city 2)
• Add 1 to ans[1]

Complete enumeration results:

• Subtrees with max distance 1: {1,2}, {2,3}, {2,4} → count = 3
• Subtrees with max distance 2: {1,2,3}, {1,2,4}, {2,3,4}, {1,2,3,4} → count = 4
• Subtrees with max distance 3: none → count = 0

Final answer: [3, 4, 0]

The algorithm efficiently identifies valid connected subtrees by ensuring all bits are toggled during DFS traversal, then uses the two-DFS technique to find the diameter of each valid subtree.

Solution Implementation

Time and Space Complexity

Time Complexity: O(2^n * n)

The algorithm iterates through all possible subsets of nodes using a bitmask from 1 to 2^n - 1, which gives us O(2^n) iterations. For each valid subset (non-power-of-2 masks), it performs two DFS traversals. Each DFS traversal visits at most n nodes and for each node, it checks all adjacent nodes in the graph. Since the graph is a tree with n-1 edges, each node has at most n-1 neighbors, but across all nodes in a single DFS, we examine each edge at most twice (once from each endpoint). Therefore, each DFS takes O(n) time. The total time complexity is O(2^n * n).

Space Complexity: O(n)

The space complexity consists of:

• The adjacency list g which stores all edges twice (once for each direction), using O(n) space since there are n-1 edges in a tree
• The result array ans of size n-1, which is O(n)
• The recursion stack for DFS which can go up to depth n in the worst case (a linear tree), contributing O(n)
• A few integer variables (mask, msk, mx, nxt, cur) which use O(1) space

The overall space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Connectivity Check After First DFS

The Pitfall: A common mistake is checking if the subset forms a connected component by verifying if the visited mask equals zero after the first DFS. However, if you forget to properly initialize or update the mask during DFS, this check becomes unreliable. Specifically, using assignment instead of XOR operation, or checking the wrong mask variable can lead to incorrect results.

Incorrect Example:

def dfs(node, distance=0):
    visited_mask = visited_mask & ~(1 << node)  # Wrong! Creates local variable
    # ... rest of DFS

Solution: Always use the nonlocal keyword to modify the outer scope variable and use XOR to toggle bits:

def dfs(node, distance=0):
    nonlocal visited_mask
    visited_mask ^= (1 << node)  # Correct: toggles the bit

2. Off-by-One Error in Diameter Calculation

The Pitfall: The diameter represents the number of edges in the longest path, but it's easy to confuse this with the number of nodes. When storing results, forgetting that the array is 0-indexed while diameters start from 1 leads to index errors or wrong counts.

Incorrect Example:

result[max_distance] += 1  # IndexError when max_distance = n-1

Solution: Always subtract 1 when indexing the result array:

result[max_distance - 1] += 1  # Correct: diameter d goes to index d-1

3. Starting DFS from a Node Not in the Subset

The Pitfall: When finding a starting node for DFS, you might accidentally pick a node that's not in the current subset, especially if using operations like bit_length() incorrectly or hardcoding a starting node.

Incorrect Example:

start_node = 0  # Wrong! Node 0 might not be in the subset
dfs(start_node)

Solution: Use bit manipulation to find any set bit in the mask:

start_node = subset_mask.bit_length() - 1  # Gets the highest set bit
# Alternative: start_node = (subset_mask & -subset_mask).bit_length() - 1  # Gets the lowest set bit

4. Not Resetting Variables Between DFS Calls

The Pitfall: The algorithm requires two DFS calls per valid subset. Forgetting to reset the visited mask and max_distance between these calls will produce incorrect diameters.

Incorrect Example:

if visited_mask == 0:
    # Forgot to reset visited_mask and max_distance!
    dfs(farthest_node)
    result[max_distance - 1] += 1

Solution: Always reset the necessary variables before the second DFS:

if visited_mask == 0:
    visited_mask = subset_mask  # Reset to original subset
    max_distance = 0            # Reset distance counter
    dfs(farthest_node)
    result[max_distance - 1] += 1

5. Misunderstanding the Tree Diameter Algorithm

The Pitfall: Some might try to find the diameter with just one DFS or by checking all pairs of nodes. The two-DFS approach is crucial: the first finds any farthest node from an arbitrary start, and the second finds the actual diameter starting from that farthest node.

Incorrect Example:

# Wrong: Single DFS doesn't guarantee finding the diameter
dfs(start_node)
if visited_mask == 0:
    result[max_distance - 1] += 1  # This max_distance might not be the diameter

Solution: Always perform two DFS traversals for finding the tree diameter:

# First DFS: find one end of the diameter
dfs(start_node)
if visited_mask == 0:  # Connected component check
    # Second DFS: find actual diameter from the farthest node
    visited_mask = subset_mask
    max_distance = 0
    dfs(farthest_node)
    result[max_distance - 1] += 1