LeetCode 1617 - Count Subtrees With Max Distance Between Cities

The input describes an undirected tree with n cities. Since the graph is a tree, there are exactly n - 1 edges and there is a unique simple path between every pair of cities. The problem asks us to examine every possible connected subset of cities.

leetcodeharddynamic-programmingbit-manipulationtreeenumerationbitmask

LeetCode Problem 1617

Difficulty: 🔴 Hard
Topics: Dynamic Programming, Bit Manipulation, Tree, Enumeration, Bitmask

Solution

Problem Understanding

The input describes an undirected tree with n cities. Since the graph is a tree, there are exactly n - 1 edges and there is a unique simple path between every pair of cities.

The problem asks us to examine every possible connected subset of cities. A subset is considered a valid subtree if all cities inside the subset remain connected using only nodes from the subset itself. In graph terms, we are looking for all connected induced subgraphs of the tree.

For every valid subtree, we compute its diameter. The diameter is the maximum distance between any two cities inside that subtree, where distance means the number of edges on the shortest path between the two cities.

The task is to count how many connected subtrees have diameter 1, how many have diameter 2, and so on up to n - 1.

The returned array has length n - 1, where:

  • answer[0] stores the number of connected subtrees with diameter 1
  • answer[1] stores the number of connected subtrees with diameter 2
  • ...
  • answer[n - 2] stores the number of connected subtrees with diameter n - 1

The constraints are extremely important:

  • 2 <= n <= 15

A tree with only 15 nodes is very small. This immediately suggests that bitmask enumeration is feasible. Since there are at most 2^15 = 32768 subsets, we can realistically iterate over every subset of nodes.

The challenge is not the number of subsets, but efficiently determining:

  1. Whether a subset forms a connected subtree
  2. What the diameter of that subtree is

Several edge cases are important:

  • A subset with only one node is connected, but its diameter is 0, and the problem only asks for diameters from 1 to n - 1, so single-node subsets should not contribute to the answer.
  • Some subsets are disconnected and must be ignored.
  • Trees can be shaped very differently, such as stars, chains, or balanced trees, so the algorithm must work for all topologies.
  • Since the graph is already a tree, every connected subset automatically forms a valid subtree.

Approaches

Brute Force Approach

The most direct solution is to enumerate every subset of nodes and test whether it forms a connected subtree.

For each subset:

  1. Extract all included nodes.
  2. Check whether the induced graph is connected using DFS or BFS.
  3. If connected, compute the maximum distance between every pair of nodes inside the subset.
  4. Increment the answer bucket corresponding to the diameter.

To compute distances, we could run BFS from every node inside the subset.

This approach is correct because it explicitly evaluates every possible subtree and directly measures its diameter. However, it performs a large amount of repeated work.

If we recompute distances repeatedly for every subset, the complexity becomes unnecessarily high.

Key Insight

The constraint n <= 15 allows full subset enumeration, but we should optimize the per-subset work.

The important observation is that the graph is a tree, meaning:

  • A connected subset with k nodes must contain exactly k - 1 internal edges.
  • The diameter of a tree can be found efficiently using two BFS or DFS traversals.

We can therefore:

  1. Enumerate all subsets using bitmasks.
  2. Check connectivity with DFS.
  3. Compute the diameter using the classic two-pass DFS technique.

The two-pass diameter method works as follows:

  • Start from any node in the subtree.
  • Find the farthest node A.
  • Start from A and find the farthest node B.
  • The distance from A to B is the subtree diameter.

This avoids checking every pair explicitly.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n * n^3) O(n^2) Enumerates subsets and checks all node pairs
Optimal O(2^n * n^2) O(n^2) Uses bitmask enumeration with DFS-based diameter computation

Algorithm Walkthrough

Step 1: Build the Adjacency List

Convert the edge list into an adjacency list representation.

This allows efficient DFS traversal of the tree.

Step 2: Enumerate All Subsets

Use a bitmask from 1 to (1 << n) - 1.

Each bit represents whether a node is included in the current subset.

For example, with n = 4:

  • 0101 means nodes {1, 3}
  • 1110 means nodes {2, 3, 4}

We skip subsets with only one node because their diameter is 0.

Step 3: Find a Starting Node

For the current subset, locate any node whose bit is set.

This node becomes the DFS starting point.

Step 4: Check Connectivity

Perform DFS while only traversing nodes included in the subset.

Track how many nodes are visited.

If the number of visited nodes is smaller than the subset size, the subset is disconnected and should be ignored.

Step 5: Find One Endpoint of the Diameter

Run DFS from the starting node.

Track the farthest reachable node and its distance.

This produces one endpoint of the subtree diameter.

Step 6: Compute the Diameter

Run DFS again, this time starting from the farthest node found in Step 5.

The maximum distance reached during this traversal is the diameter.

Step 7: Update the Answer

If the diameter is greater than 0, increment:

answer[diameter - 1]

because the result array is zero-indexed.

Step 8: Continue Until All Subsets Are Processed

After examining every subset, return the answer array.

Why it works

Every possible subset of nodes is examined exactly once. A subset contributes to the answer only if it is connected, which guarantees it forms a valid subtree in the original tree.

The two-pass DFS method correctly computes the diameter of a tree because one endpoint of the diameter is always the farthest node from any arbitrary starting node. Running DFS again from that endpoint yields the true maximum distance.

Since every connected subtree is counted exactly once and assigned to its correct diameter bucket, the algorithm is correct.

Python Solution

from typing import List

class Solution:
    def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = [[] for _ in range(n)]

        for u, v in edges:
            u -= 1
            v -= 1
            graph[u].append(v)
            graph[v].append(u)

        answer = [0] * (n - 1)

        def dfs(node: int, mask: int, visited: List[bool], distance: int):
            visited[node] = True

            farthest_node = node
            max_distance = distance
            count = 1

            for neighbor in graph[node]:
                if ((mask >> neighbor) & 1) and not visited[neighbor]:
                    next_node, next_distance, next_count = dfs(
                        neighbor,
                        mask,
                        visited,
                        distance + 1
                    )

                    count += next_count

                    if next_distance > max_distance:
                        max_distance = next_distance
                        farthest_node = next_node

            return farthest_node, max_distance, count

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

            start = 0
            while ((mask >> start) & 1) == 0:
                start += 1

            visited = [False] * n

            farthest_node, _, visited_count = dfs(
                start,
                mask,
                visited,
                0
            )

            subset_size = mask.bit_count()

            if visited_count != subset_size:
                continue

            visited = [False] * n

            _, diameter, _ = dfs(
                farthest_node,
                mask,
                visited,
                0
            )

            answer[diameter - 1] += 1

        return answer

The implementation begins by converting the edge list into an adjacency list. Since the input nodes are one-indexed but Python lists are zero-indexed, each node value is decremented by one.

The core helper function is dfs. It performs three tasks simultaneously:

  1. Traverses only nodes inside the current subset
  2. Counts how many nodes are visited
  3. Tracks the farthest reachable node and its distance

The recursion returns three values:

  • The farthest node found
  • The distance to that node
  • The number of visited nodes

The outer loop iterates through every bitmask subset.

Single-node subsets are skipped immediately because they do not contribute to the answer.

The first DFS checks connectivity and also finds one endpoint candidate for the diameter. If the visited count does not equal the number of nodes in the subset, the subset is disconnected.

The second DFS starts from the farthest node found earlier. The maximum distance discovered in this traversal is the diameter.

Finally, the corresponding answer bucket is incremented.

Go Solution

package main

func countSubgraphsForEachDiameter(n int, edges [][]int) []int {
	graph := make([][]int, n)

	for _, edge := range edges {
		u := edge[0] - 1
		v := edge[1] - 1

		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	answer := make([]int, n-1)

	var dfs func(int, int, []bool, int) (int, int, int)

	dfs = func(node int, mask int, visited []bool, distance int) (int, int, int) {
		visited[node] = true

		farthestNode := node
		maxDistance := distance
		count := 1

		for _, neighbor := range graph[node] {
			if ((mask>>neighbor)&1) == 1 && !visited[neighbor] {
				nextNode, nextDistance, nextCount := dfs(
					neighbor,
					mask,
					visited,
					distance+1,
				)

				count += nextCount

				if nextDistance > maxDistance {
					maxDistance = nextDistance
					farthestNode = nextNode
				}
			}
		}

		return farthestNode, maxDistance, count
	}

	for mask := 1; mask < (1 << n); mask++ {
		if mask&(mask-1) == 0 {
			continue
		}

		start := 0

		for ((mask >> start) & 1) == 0 {
			start++
		}

		visited := make([]bool, n)

		farthestNode, _, visitedCount := dfs(
			start,
			mask,
			visited,
			0,
		)

		subsetSize := 0

		for temp := mask; temp > 0; temp >>= 1 {
			subsetSize += temp & 1
		}

		if visitedCount != subsetSize {
			continue
		}

		visited = make([]bool, n)

		_, diameter, _ := dfs(
			farthestNode,
			mask,
			visited,
			0,
		)

		answer[diameter-1]++
	}

	return answer
}

The Go implementation follows the same structure as the Python solution.

One notable difference is that Go does not provide a built-in bit_count() function for integers in the same convenient way as Python. Therefore, the subset size is computed manually by counting set bits.

Slices are used for adjacency lists and visited tracking. Since Go passes slices by reference-like behavior, the DFS function can modify the shared visited slice directly.

Integer overflow is not a concern because the maximum values involved are extremely small due to the constraint n <= 15.

Worked Examples

Example 1

Input:

n = 4
edges = [[1,2],[2,3],[2,4]]

The tree structure is:

    1
    |
    2
   / \
  3   4

We enumerate all subsets.

Subset Connected Diameter
{1} Yes 0
{2} Yes 0
{3} Yes 0
{4} Yes 0
{1,2} Yes 1
{2,3} Yes 1
{2,4} Yes 1
{1,3} No Invalid
{1,4} No Invalid
{3,4} No Invalid
{1,2,3} Yes 2
{1,2,4} Yes 2
{2,3,4} Yes 2
{1,2,3,4} Yes 2

Final counts:

  • Diameter 1: 3
  • Diameter 2: 4
  • Diameter 3: 0

Output:

[3,4,0]

Example 2

Input:

n = 2
edges = [[1,2]]

Possible subsets:

Subset Connected Diameter
{1} Yes 0
{2} Yes 0
{1,2} Yes 1

Output:

[1]

Example 3

Input:

n = 3
edges = [[1,2],[2,3]]

Tree structure:

1 - 2 - 3

Subsets:

Subset Connected Diameter
{1,2} Yes 1
{2,3} Yes 1
{1,3} No Invalid
{1,2,3} Yes 2

Output:

[2,1]

Complexity Analysis

Measure Complexity Explanation
Time O(2^n * n^2) Each subset may require DFS traversals across up to n nodes
Space O(n^2) Adjacency list plus recursion stack and visited arrays

The algorithm enumerates all 2^n subsets. For each subset, we may perform two DFS traversals, each taking O(n) time. Since subset operations and traversal overhead also involve node checks, the total complexity is O(2^n * n^2).

This is efficient enough because n <= 15, so the maximum number of subsets is only 32768.

Test Cases

from typing import List

class Solution:
    def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = [[] for _ in range(n)]

        for u, v in edges:
            u -= 1
            v -= 1
            graph[u].append(v)
            graph[v].append(u)

        answer = [0] * (n - 1)

        def dfs(node: int, mask: int, visited: List[bool], distance: int):
            visited[node] = True

            farthest_node = node
            max_distance = distance
            count = 1

            for neighbor in graph[node]:
                if ((mask >> neighbor) & 1) and not visited[neighbor]:
                    next_node, next_distance, next_count = dfs(
                        neighbor,
                        mask,
                        visited,
                        distance + 1
                    )

                    count += next_count

                    if next_distance > max_distance:
                        max_distance = next_distance
                        farthest_node = next_node

            return farthest_node, max_distance, count

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

            start = 0
            while ((mask >> start) & 1) == 0:
                start += 1

            visited = [False] * n

            farthest_node, _, visited_count = dfs(
                start,
                mask,
                visited,
                0
            )

            if visited_count != mask.bit_count():
                continue

            visited = [False] * n

            _, diameter, _ = dfs(
                farthest_node,
                mask,
                visited,
                0
            )

            answer[diameter - 1] += 1

        return answer

solution = Solution()

assert solution.countSubgraphsForEachDiameter(
    4,
    [[1,2],[2,3],[2,4]]
) == [3,4,0]  # official example 1

assert solution.countSubgraphsForEachDiameter(
    2,
    [[1,2]]
) == [1]  # smallest valid tree

assert solution.countSubgraphsForEachDiameter(
    3,
    [[1,2],[2,3]]
) == [2,1]  # chain structure

assert solution.countSubgraphsForEachDiameter(
    3,
    [[1,2],[1,3]]
) == [2,1]  # star of size 3

assert solution.countSubgraphsForEachDiameter(
    5,
    [[1,2],[2,3],[3,4],[4,5]]
) == [4,3,2,1]  # long chain

assert solution.countSubgraphsForEachDiameter(
    5,
    [[1,2],[1,3],[1,4],[1,5]]
) == [4,11,0,0]  # star graph

assert solution.countSubgraphsForEachDiameter(
    6,
    [[1,2],[2,3],[2,4],[4,5],[4,6]]
) == [5,8,9,0,0]  # branching tree

Test Summary

Test Why
n=4, branching tree Verifies official example
n=2 Smallest valid input
Chain of 3 nodes Tests simple diameter growth
Star with 3 nodes Verifies disconnected subset handling
Chain of 5 nodes Tests increasing diameters
Star of 5 nodes Tests many diameter-2 subtrees
Larger branching tree Stress test for mixed structures

Edge Cases

A very important edge case is single-node subsets. Technically, a single node forms a connected subtree, but its diameter is 0. Since the problem only asks for diameters from 1 to n - 1, these subsets must not contribute to the result. The implementation handles this by skipping masks that contain only one set bit.

Another important edge case is disconnected subsets. For example, in a chain 1-2-3, the subset {1,3} is invalid because node 2 is missing, so the remaining nodes are disconnected. A naive implementation might accidentally treat this as a valid subtree. The DFS connectivity check prevents this by ensuring the number of visited nodes matches the subset size.

Star-shaped trees are also tricky. In a star centered at node 1, any subset containing two leaf nodes must also include the center node to remain connected. The implementation naturally enforces this because DFS traversal cannot reach between leaves without the center.

Finally, long chain trees stress the diameter computation. In a path graph, the diameter can become as large as n - 1. The two-pass DFS method correctly identifies the endpoints of the longest path regardless of tree shape.