Tree Diameter

Intuition

For all nodes in the diameter, only one node is allowed that both left edge and right edge are used. So we could use a DFS to get this diameter.

Solution


from collections import defaultdict


class Solution:
    def treeDiameter(self, edges: List[List[int]]) -> int:
        if not edges:
            return 0

        adj = defaultdict(list)
        for p, q in edges:
            adj[p].append(q)
            adj[q].append(p)
        visited = set()
        root = edges[0][0]
        visited.add(root)
        return max(self._helper(root, adj, visited))

    def _helper(self, node, adj, visited):
        visited.add(node)
        m1, m2, closed = 0, 0, 0
        for child in adj[node]:
            if child in visited:
                continue
            m, m_closed = self._helper(child, adj, visited)
            m += 1
            _, m1, m2 = sorted([m1, m2, m])
            closed = max(closed, m_closed)
        closed = max(closed, m1 + m2)
        return max(m1, m2), closed