Skip to content

1245. Tree Diameter πŸ”’

SourceBiweekly Contest 12 Q3DifficultyMediumRating1792

Description

The diameter of a tree is the number of edges in the longest path in that tree.

There is an undirected tree of n nodes labeled from 0 to n - 1. You are given a 2D array edges where edges.length == n - 1 and edges[i] = [ai, bi] indicates that there is an undirected edge between nodes ai and bi in the tree.

Return the diameter of the tree.

 

Example 1:

Input: edges = [[0,1],[0,2]]
Output: 2
Explanation: The longest path of the tree is the path 1 - 0 - 2.

Example 2:

Input: edges = [[0,1],[1,2],[2,3],[1,4],[4,5]]
Output: 4
Explanation: The longest path of the tree is the path 3 - 2 - 1 - 4 - 5.

 

Constraints:

  • n == edges.length + 1
  • 1 <= n <= 104
  • 0 <= ai, bi < n
  • ai != bi

Solutions

Solution 1: Two DFS Passes

Thinking

The diameter is the longest simple path. \(n \le 10^4\) forbids all endpoint pairs. From any vertex the farthest vertex is an end of some diameter; a second walk from that end is the diameter.

Two DFS passes: the first from \(0\) finds the farthest \(a\); the second from \(a\) is the length. The graph is a tree, so DFS and BFS both find a farthest vertex.

First, we arbitrarily select a node and start a depth-first search (DFS) from this node to find the farthest node from it, denoted as node \(a\). Then, we start another DFS from node \(a\) to find the farthest node from node \(a\), denoted as node \(b\). It can be proven that the path between node \(a\) and node \(b\) is the diameter of the tree.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the number of nodes.

Similar problems:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def treeDiameter(self, edges: list[list[int]]) -> int:
        def dfs(i: int, fa: int, t: int):
            for j in g[i]:
                if j != fa:
                    dfs(j, i, t + 1)
            nonlocal ans, a
            if ans < t:
                ans = t
                a = i

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = a = 0
        dfs(0, -1, 0)
        dfs(a, -1, 0)
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution {
    private List<Integer>[] g;
    private int ans;
    private int a;

    public int treeDiameter(int[][] edges) {
        int n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        dfs(0, -1, 0);
        dfs(a, -1, 0);
        return ans;
    }

    private void dfs(int i, int fa, int t) {
        for (int j : g[i]) {
            if (j != fa) {
                dfs(j, i, t + 1);
            }
        }
        if (ans < t) {
            ans = t;
            a = i;
        }
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution {
public:
    int treeDiameter(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<int> g[n];
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        int ans = 0, a = 0;
        auto dfs = [&](this auto&& dfs, int i, int fa, int t) -> void {
            for (int j : g[i]) {
                if (j != fa) {
                    dfs(j, i, t + 1);
                }
            }
            if (ans < t) {
                ans = t;
                a = i;
            }
        };
        dfs(0, -1, 0);
        dfs(a, -1, 0);
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
func treeDiameter(edges [][]int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    a := 0
    var dfs func(i, fa, t int)
    dfs = func(i, fa, t int) {
        for _, j := range g[i] {
            if j != fa {
                dfs(j, i, t+1)
            }
        }
        if ans < t {
            ans = t
            a = i
        }
    }
    dfs(0, -1, 0)
    dfs(a, -1, 0)
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
function treeDiameter(edges: number[][]): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    let [ans, a] = [0, 0];
    const dfs = (i: number, fa: number, t: number): void => {
        for (const j of g[i]) {
            if (j !== fa) {
                dfs(j, i, t + 1);
            }
        }
        if (ans < t) {
            ans = t;
            a = i;
        }
    };
    dfs(0, -1, 0);
    dfs(a, -1, 0);
    return ans;
}

Solution 2: Two Explicit Stack Passes

Thinking

The diameter is the number of edges on a longest simple path. With \(n \le 10^4\), trying every pair of endpoints is too slow. The farthest vertex from any start is an endpoint of some diameter, and a second walk from that endpoint yields the length. Recursing along a chain uses a call depth equal to the node count and overflows Python once the chain reaches length \(1000\). Each search therefore stores \((\textit{node}, \textit{parent}, \textit{distance})\) on a stack: a larger distance on pop records the current node, then the other neighbors are pushed at distance one greater. The first search from \(0\) finds endpoint \(a\), and the farthest distance from \(a\) is the diameter.

We start at node \(0\) and use an explicit stack to find the farthest node \(a\), then repeat the search from \(a\). The farthest distance in the second search is the diameter of the tree.

Each stack frame is \((i, fa, t)\): the walk has reached node \(i\) from parent \(fa\) after \(t\) edges. When the frame is popped, a larger \(t\) replaces the farthest node and the answer, and every neighbor other than the parent is pushed at distance \(t + 1\). The tree has no cycle, so each node is pushed once. Which tied endpoint is recorded does not change the length found by the second search.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the number of nodes.

Similar problems:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution:
    def treeDiameter(self, edges: list[list[int]]) -> int:
        n = len(edges) + 1
        g = [[] for _ in range(n)]
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)

        def farthest(start: int) -> (int, int):
            ans = 0
            node = start
            stk = [(start, -1, 0)]
            while stk:
                i, fa, t = stk.pop()
                if ans < t:
                    ans = t
                    node = i
                for j in g[i]:
                    if j != fa:
                        stk.append((j, i, t + 1))
            return node, ans

        node, _ = farthest(0)
        return farthest(node)[1]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
    public int treeDiameter(int[][] edges) {
        int n = edges.length + 1;
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        int node = farthest(g, 0)[0];
        return farthest(g, node)[1];
    }

    private int[] farthest(List<Integer>[] g, int start) {
        int ans = 0, node = start;
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {start, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], t = cur[2];
            if (ans < t) {
                ans = t;
                node = i;
            }
            for (int j : g[i]) {
                if (j != fa) {
                    stk.push(new int[] {j, i, t + 1});
                }
            }
        }
        return new int[] {node, ans};
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution {
public:
    int treeDiameter(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        auto farthest = [&](int start) -> pair<int, int> {
            int ans = 0, node = start;
            vector<array<int, 3>> stk{{start, -1, 0}};
            while (!stk.empty()) {
                auto [i, fa, t] = stk.back();
                stk.pop_back();
                if (ans < t) {
                    ans = t;
                    node = i;
                }
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push_back({j, i, t + 1});
                    }
                }
            }
            return {node, ans};
        };
        int node = farthest(0).first;
        return farthest(node).second;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
func treeDiameter(edges [][]int) int {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    farthest := func(start int) (int, int) {
        ans, node := 0, start
        stk := [][3]int{{start, -1, 0}}
        for len(stk) > 0 {
            cur := stk[len(stk)-1]
            stk = stk[:len(stk)-1]
            i, fa, t := cur[0], cur[1], cur[2]
            if ans < t {
                ans, node = t, i
            }
            for _, j := range g[i] {
                if j != fa {
                    stk = append(stk, [3]int{j, i, t + 1})
                }
            }
        }
        return node, ans
    }
    node, _ := farthest(0)
    _, ans := farthest(node)
    return ans
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
function treeDiameter(edges: number[][]): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    const farthest = (start: number): [number, number] => {
        let ans = 0;
        let node = start;
        const stk: number[][] = [[start, -1, 0]];
        while (stk.length) {
            const [i, fa, t] = stk.pop()!;
            if (ans < t) {
                ans = t;
                node = i;
            }
            for (const j of g[i]) {
                if (j !== fa) {
                    stk.push([j, i, t + 1]);
                }
            }
        }
        return [node, ans];
    };
    const [node] = farthest(0);
    return farthest(node)[1];
}

Comments