1245. Tree Diameter


Problem Description

The task is to find the diameter of a tree. The diameter is defined as the number of edges in the longest path between any two nodes in the tree. Given that the tree is undirected and consists of n nodes labeled from 0 to n - 1, and an array of edges indicating the connections between the nodes, the objective is to calculate and return the tree's diameter.

Intuition

Finding the diameter of a tree involves discovering the longest path between two nodes. A common technique to solve this problem is to use a depth-first search (DFS) twice. The intuition behind this approach is based on the property that the longest path in the tree will have as its endpoints two of the tree's leaves.

The first DFS begins at an arbitrary node (usually the first node provided) and traverses the tree to find the leaf node that lies at the greatest distance from the starting node. This leaf node will be one end of the longest path. The second DFS then begins at this leaf node and searches for the farthest leaf node from it, thus uncovering the other end of the longest path.

By keeping track of the number of edges traversed during the second DFS (which is essentially the length or the height of the tree from the perspective of the starting leaf node), we can directly calculate the diameter of the tree.

Learn more about Tree, Depth-First Search, Breadth-First Search, Graph and Topological Sort patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

Suppose k is a very large integer(2^64). Which of the following is the largest as n grows to infinity?

Solution Approach

The provided Python code implements a solution to the problem using a depth-first search (DFS) algorithm. Here's how it works:

  1. A dfs function is defined, which takes a node u and a counter t that tracks the total number of edges traversed from the starting node. The dfs function recursively visits nodes, incrementing the counter as it goes deeper into the tree.

  2. Inside this DFS function, a base condition checks if the current node has already been visited to prevent loops. If it hasn't been visited, it is marked as visited.

  3. The recursive step involves iterating over all the adjacent nodes (v) of the current node (u). For each adjacent node, the dfs function is called again with v as the current node and the counter t incremented by 1.

  4. While returning from each call of dfs, the code compares the current value of t with a variable ans, which keeps track of the maximum distance found so far. If t is greater, it means a longer path has been found, and thus ans is updated with the current value of t, and next is updated with the current node u.

The main function treeDiameter utilizes the DFS as follows:

  • A defaultdict d is created to store the tree in an adjacency list representation. This makes it easy to access all the nodes connected to a given node in the tree. The edges are used to populate this adjacency list.

  • An array vis is used to keep track of which nodes have been visited to prevent repeating the same nodes during the DFS.

  • The first DFS is initiated from the first node in the edges list. After this DFS, we have a candidate leaf node for the longest path in the variable next.

  • The vis array is reset to prepare for the second DFS call.

  • The second DFS is initiated from the next node found from the first DFS call.

  • At the end of the second DFS, the variable ans holds the diameter of the tree, which is the number of edges in the longest path. This is returned as the function's result.

In summary, the algorithm finds the longest path (or diameter) of the tree by first locating one end of the longest path using a DFS and then executing another DFS from that endpoint to find the actual longest path.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

What are the most two important steps in writing a depth first search function? (Select 2)

Example Walkthrough

Let's consider a small example to illustrate the solution approach. Suppose we have a tree with n=6 nodes and the following edges:

10 - 1
21 - 2
31 - 3
42 - 4
52 - 5

The tree structure would look like this:

1  0
2   \
3    1
4   / \
5  2   3
6 / \
74   5

Given this tree, we want to find its diameter. Following the solution approach:

  1. We first use DFS starting from an arbitrary node, say node 0.

  2. We traverse the tree from node 0, and we may reach a leaf node, such as node 4. Let's assume the DFS process went along the path 0->1->2->4. During this process, each time DFS calls itself, it increments the counter t.

  3. Since node 4 is a leaf, we backtrack, but in doing so, we keep track of the node that gave us the maximum depth. Let's assume this maximum depth was first encountered at node 4.

  4. With node 4 as our furthest leaf from node 0, we now reset the vis array for the second DFS.

  5. The second DFS starts at node 4 and proceeds to find the furthest node from it. Node 4 will branch only to 2, so we go to 2, then to 1, and we have a choice to go to node 0 or node 3.

  6. If we go to 3, which is another leaf, we've found the longest path in the tree: 4->2->1->3.

During the second DFS, we keep a variable ans that tracks the max length. Since the path 4->2->1->3 is the longest path in the tree, ans would be 3 in the end, which means there are three edges in the path. Thus, the diameter of this tree is 3.

That concludes the example walkthrough using the solution approach. We've illustrated how the algorithm recursively searches for the furthest possible nodes step by step and eventually finds the longest path, which defines the diameter of the tree.

Not Sure What to Study? Take the 2-min Quiz:

Breadth first search can be used to find the shortest path between two nodes in a directed graph.

Python Solution

1from collections import defaultdict
2
3class Solution:
4    # Function to find the diameter of a tree given its edges
5    def treeDiameter(self, edges: List[List[int]]) -> int:
6        # Helper function to perform a depth-first search (DFS)
7        def dfs(node, distance):
8            nonlocal max_diameter, visited, graph, next_node
9            if visited[node]:  # If this node has already been visited, skip
10                return
11            visited[node] = True
12            for neighbor in graph[node]:
13                dfs(neighbor, distance + 1)  # DFS on connected nodes
14          
15            # Update the diameter and the next node if we found a longer path
16            if max_diameter < distance:
17                max_diameter = distance
18                next_node = node
19      
20        # Convert the edge-list to an adjacency list representation for the graph
21        graph = defaultdict(set)
22        # Initialize a visited list to keep track of visited nodes during DFS
23        visited = [False] * (len(edges) + 1)
24        # Establish the adjacency list from the edges
25        for u, v in edges:
26            graph[u].add(v)
27            graph[v].add(u)
28      
29        max_diameter = 0  # Initialize the diameter of the tree
30        next_node = 0  # This variable will hold one end of the maximum diameter
31      
32        # Perform the first DFS to find one end of the maximum diameter path
33        dfs(edges[0][0], 0)
34      
35        # Reset the visited list for the next DFS
36        visited = [False] * (len(edges) + 1)
37      
38        # Perform the second DFS from the farthest node found in the first DFS
39        dfs(next_node, 0)
40
41        # The maximum distance found during the second DFS is the tree diameter
42        return max_diameter
43

Java Solution

1class Solution {
2    private Map<Integer, Set<Integer>> graph; // Represents the adjacency list for the graph
3    private boolean[] visited;                 // Array to keep track of visited nodes
4    private int farthestNode;                  // Tracks the farthest node found during DFS
5    private int diameter;                      // Stores the current diameter of the tree
6
7    public int treeDiameter(int[][] edges) {
8        int n = edges.length;       // Number of nodes will be #edges + 1 in a tree
9        diameter = 0;               // Initialize diameter as zero
10        graph = new HashMap<>();    // Initialize the graph
11
12        // Building an undirected graph using the provided edges
13        for (int[] edge : edges) {
14            graph.computeIfAbsent(edge[0], k -> new HashSet<>()).add(edge[1]);
15            graph.computeIfAbsent(edge[1], k -> new HashSet<>()).add(edge[0]);
16        }
17
18        // Initialize the visited array for the first DFS traversal
19        visited = new boolean[n + 1]; // Index 0 is ignored since node numbering starts from 1
20        farthestNode = edges[0][0];   // Start from the first node
21        dfs(farthestNode, 0);         // Perform DFS to find the farthest node
22
23        // Reset the visited array for the second DFS traversal
24        visited = new boolean[n + 1];
25        dfs(farthestNode, 0);         // Perform DFS from the farthest node to find the diameter
26        return diameter;              // Return the diameter
27    }
28
29    private void dfs(int node, int distance) {
30        if (visited[node]) {
31            return;
32        }
33        visited[node] = true;
34        // Update the diameter and farthest node if the current distance is greater
35        if (diameter < distance) {
36            diameter = distance;
37            farthestNode = node;
38        }
39        // Visit all the connected nodes of the current node using DFS
40        for (int adjacent : graph.get(node)) {
41            dfs(adjacent, distance + 1);
42        }
43    }
44}
45

C++ Solution

1#include <vector>
2#include <unordered_map>
3#include <unordered_set>
4
5class Solution {
6public:
7    // Graph representation using an adjacency list
8    unordered_map<int, unordered_set<int>> graph;
9    // Visited vector to keep track of visited nodes during DFS
10    vector<bool> visited;
11    // Variable to store the length of the diameter of the tree
12    int diameter;
13    // Variable to store the node that will be the next candidate for DFS
14    int furthestNode;
15
16    // Function to calculate the diameter of the tree
17    int treeDiameter(vector<vector<int>>& edges) {
18        // Construct the graph
19        for (auto& edge : edges) {
20            graph[edge[0]].insert(edge[1]);
21            graph[edge[1]].insert(edge[0]);
22        }
23        int numNodes = edges.size();
24        diameter = 0;
25        // Initialize visited vector of size numNodes+1 (since it's a tree, it has numNodes+1 vertices)
26        visited.resize(numNodes + 1, false);
27        // Initialize the furthestNode with any node (e.g., the first node of the first edge)
28        furthestNode = edges[0][0];
29      
30        // Perform the first DFS to find the furthest node from the starting node
31        dfs(furthestNode, 0);
32        // Reset the visited vector for the next DFS
33        visited.assign(visited.size(), false);
34        // Perform the DFS starting from the furthest node to determine the diameter
35        dfs(furthestNode, 0);
36      
37        // Diameter is the maximum number of edges between any two nodes in the tree
38        return diameter;
39    }
40
41    // Depth-First Search (DFS) function to traverse the tree
42    void dfs(int currentNode, int currentLength) {
43        // If the current node has been visited, return
44        if (visited[currentNode]) return;
45        // Mark the current node as visited
46        visited[currentNode] = true;
47        // If the current length is greater than the diameter, update the diameter and furthestNode
48        if (diameter < currentLength) {
49            diameter = currentLength;
50            furthestNode = currentNode;
51        }
52        // Visit all adjacent nodes
53        for (int adjacentNode : graph[currentNode]) {
54            dfs(adjacentNode, currentLength + 1);
55        }
56    }
57};
58

Typescript Solution

1// Importing necessary functionalities from 'collections' module in TypeScript
2import { Set, Map } from "typescript-collections";
3
4// Global graph representation using an adjacency list
5const graph: Map<number, Set<number>> = new Map();
6// Global visited Map to keep track of visited nodes during DFS
7const visited: Map<number, boolean> = new Map();
8// Variable to store the length of the diameter of the tree
9let diameter: number = 0;
10// Variable to store the node that will be the next candidate for DFS
11let furthestNode: number;
12
13// Function to calculate the diameter of the tree
14function treeDiameter(edges: number[][]): number {
15    // Construct the graph
16    edges.forEach(edge => {
17        if (!graph.containsKey(edge[0])) {
18            graph.setValue(edge[0], new Set());
19        }
20        if (!graph.containsKey(edge[1])) {
21            graph.setValue(edge[1], new Set());
22        }
23        graph.getValue(edge[0]).add(edge[1]);
24        graph.getValue(edge[1]).add(edge[0]);
25    });
26
27    const numNodes: number = edges.length;
28    diameter = 0;
29    // Initialize visited Map for all nodes as false
30    for (let i = 0; i < numNodes + 1; i++) {
31        visited.setValue(i, false);
32    }
33    // Initialize the furthestNode with any node (e.g., the first node of the first edge)
34    furthestNode = edges[0][0];
35
36    // Perform the first DFS to find the furthest node from the starting node
37    dfs(furthestNode, 0);
38    // Reset the visited Map for the next DFS
39    visited.forEach((value, key) => {
40        visited.setValue(key, false);
41    });
42    // Perform the DFS starting from the furthest node to determine the diameter
43    dfs(furthestNode, 0);
44
45    // Diameter is the maximum number of edges between any two nodes in the tree
46    return diameter;
47}
48
49// Depth-First Search (DFS) function to traverse the tree
50function dfs(currentNode: number, currentLength: number): void {
51    // If the current node has been visited, return
52    if (visited.getValue(currentNode)) {
53        return;
54    }
55    // Mark the current node as visited
56    visited.setValue(currentNode, true);
57    // If the current length is greater than the diameter, update the diameter and furthestNode
58    if (currentLength > diameter) {
59        diameter = currentLength;
60        furthestNode = currentNode;
61    }
62    // Visit all adjacent nodes
63    graph.getValue(currentNode).forEach(adjacentNode => {
64        dfs(adjacentNode, currentLength + 1);
65    });
66}
67
Fast Track Your Learning with Our Quick Skills Quiz:

Which data structure is used to implement priority queue?

Time and Space Complexity

Time Complexity

The given Python code snippet defines a Solution class with a method treeDiameter to find the diameter of an undirected tree represented by a list of edges. The algorithm uses a depth-first search (dfs) starting from an arbitrary node to find the farthest node from it, and then performs a second depth-first search from this newly discovered node to find the diameter of the tree.

Let n be the number of nodes in the tree, which is one more than the number of edges (len(edges)), as it is a connected, undirected tree with n-1 edges.

The time complexity for constructing the adjacency list d is O(n), since each edge (u, v) is processed exactly twice, once for each direction.

Each dfs call traverses all edges in the tree exactly once. Since there are n-1 edges and each edge contributes to two entries in the adjacency list (one for each end of the edge), the total number of adjacency entries to traverse is 2(n-1), which is O(n). There are two dfs calls, meaning that the total time for traversing the tree is 2*O(n), which simplifies to O(n).

Therefore, the overall time complexity of the code is O(n).

Space Complexity

The space complexity consists of the storage for the adjacency list d, the visited flags array vis, and the auxiliary space used by the dfs recursion call stack.

The adjacency list d has 2(n-1) entries as explained earlier; thus, its space complexity is O(n).

The vis array is of length n, so it also occupies O(n) space.

The depth of the recursive dfs call stack will be at most n in the worst case (a path graph), contributing an additional O(n) to the space complexity.

Combining these factors, the overall space complexity of the algorithm is O(n).

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


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫