Leetcode 1245. Tree Diameter

Problem Explanation

In this problem, we are given a tree. Our task is to compute the diameter of the tree. The diameter of a tree is defined as the "the number of edges in a longest path in that tree". The tree itself is structured as an adjacent list. An edge is a tree connects two nodes (u,v).

Solution Approach

The solution uses depth-first search (DFS) along with recursion and memorization to obtain the diameter of the tree. The 'maxDepth' function returns the maximum depth in the form of the longest path that starts at node 'u' and goes only downward. For every non-parent child 'v' of 'u', there can be a longest path that passes 'u'. The length of the path is the sum of the path from 'u' to 'v' and the maximum distance from 'u' to all the other children. We keep track of the maximum depth and the second maximum depth for any child so that we can calculate the diameter.

Let's take an example,

Input: edges = [[0,1],[1,2],[2,3],[1,4],[4,5]]

Example Walkthrough

Here, the tree looks like this:

0
1 /
2 4 | | 3 5

By executing the code, the max depth for node 0 will be 1, for node 1 it will be 3, for node 2 it will be 2, for node 3 it will 1, for node 4 it will be 2 and for node 5 it will 1. The maximum among these is the node 1 which is 3 and that will be our output.

Python Solution

1
2python
3from collections import defaultdict
4class Solution:
5    def treeDiameter(self, edges):
6        def dfs(node, parent):
7            max1 = max2 = 0
8            for child in graph[node]:
9                if child == parent: continue
10                depth = dfs(child, node)
11                if depth > max1:
12                    max1, max2 = depth, max1
13                elif depth > max2:
14                    max2 = depth
15            self.diameter = max(self.diameter, max1 + max2)
16            return max1 + 1
17
18        graph = defaultdict(list)
19        for u, v in edges:
20            graph[u].append(v)
21            graph[v].append(u)
22        self.diameter = 0
23        dfs(0, None)
24        return self.diameter

Java Solution

1
2java
3import java.util.ArrayList;
4public class Solution {
5    ArrayList<Integer>[] graph;
6    int diameter;
7    public int treeDiameter(int[][] edges) {
8        int n = edges.length;
9        graph = new ArrayList[n + 1];
10        diameter = 0;
11        for (int i = 0; i < n + 1; i++){
12            graph[i] = new ArrayList<>();
13        }
14        for (int[] edge : edges){
15            graph[edge[0]].add(edge[1]);
16            graph[edge[1]].add(edge[0]);
17        }
18        maxDepth( 0, -1);
19        return diameter;
20    }
21
22    private int maxDepth(int node, int from) {
23        int max1 = 0;
24        int max2 = 0;
25        for (int to : graph[node]) {
26            if (to == from)
27                continue;
28            int depth = maxDepth(to, node);
29            if (depth > max1) {
30                max2 = max1;
31                max1 = depth;
32            } else if (depth > max2) {
33                max2 = depth;
34            }
35        }
36
37        diameter = Math.max(diameter, max1 + max2);
38        return 1 + max1;
39    }
40}

C# Solution

1
2csharp
3public class Solution {
4    List<int>[] graph;
5    int diameter = 0;
6    public int TreeDiameter(int[][] edges) {
7        int n = edges.Length;
8        graph = new List<int>[n + 1];
9        for (int i = 0; i <= n; i++){ 
10            graph[i] = new List<int>();
11        }
12        for(int[] edge: edges){
13            graph[edge[0]].Add(edge[1]);
14            graph[edge[1]].Add(edge[0]);
15        }
16        MaxDepth(0, -1);
17        return diameter;
18    }
19    int MaxDepth(int v, int from){
20        int max1 = 0; 
21        int max2 = 0; 
22        foreach (int to in graph[v]){
23            if (to == from){
24                continue;
25            }
26            int depth = MaxDepth(to, v);
27            if (depth > max1){
28                max2 = max1;
29                max1 = depth;
30            }else if (depth > max2){
31                max2 = depth;
32            }
33        }
34        diameter = Math.Max(diameter, max1 + max2);
35        return max1 + 1;
36    }
37}

JavaScript Solution

The JavaScript solution follows the same approach as described earlier. Instead of using the hashmap approach followed in other languages, it uses arrays to keep track of the edges and the graphs. Every single edge is enumerated into two parts and pushed into the graph. Depth-First Search is implemented recursively with two variables keeping track of the maximum and second maximum values, which are initialized to zero.

1
2javascript
3class Solution {
4    treeDiameter(edges) {
5        let graph = Array(edges.length+1).fill(0).map(v=>[]);
6        let diameter = 0;
7
8        function DFS(node, parent) {
9            let max1 = 0;
10            let max2 = 0;
11            for(let v of graph[node]) {
12                if(v === parent) continue;
13                let pos = DFS(v, node);
14                if(pos > max1) {
15                    [max1, max2] = [pos, max1];
16                }
17                else if(pos > max2) {
18                    max2 = pos;
19                }
20            }
21            diameter = Math.max(diameter, max1+max2);
22            return max1 + 1;
23        }
24
25        for(let edge of edges) {
26            let u = edge[0];
27            let v = edge[1];
28            graph[u].push(v);
29            graph[v].push(u);
30        }
31        DFS(0, null);
32        return diameter;
33    }
34}
35
36let sol = new Solution();
37console.log(sol.treeDiameter([[0,1],[1,2],[2,3],[1,4],[4,5]]));

This code creates a class named Solution and defines a method treeDiameter(), which is called on an object of this class. The function takes the list of edges as input and returns the diameter of the tree. The last two lines of the code create a new object of the Solution class and call the function treeDiameter(), passing the list of edges as an argument.

To sum up, in this article, we discussed an approach to solve the problem of finding the diameter of a tree using depth-first search in Python, Java, and JavaScript languages. The main idea is to find the longest path in the tree which represents the diameter of the tree. For this, we implemented the depth-first search function and used the concept of recursion. The solutions discussed above were optimized using memorization.