https://leetcode.com/problems/minimum-height-trees

Actually, what we need to find in this problem is diameter of our tree: two nodes, that have the biggest distance between them. How we can do it? There is two-stage dfs algorighm to do it:

  1. Start from any node and run dfs, which calculate distances for all other nodes from this node as well as create parent for each node. What our fucntion will return in the end is the index of farthest node (there can be more than one and it is OK for us).
  2. Now, we need to start our dfs all over again from this node (that is why we use dfs_helper function - we need to clean our distances and parents arrays. When we run our dfs second time, we again find the farthest node.
  3. Finally, we create path between the first node we found and the second and it will be our diameter: if it has even number of nodes, we return 2 middle nodes, in other case we return 1 middle node.

Complexity: time complexity is O(n), because we use our dfs twice. Space complexity is O(n) as well.

class Solution:
    def findMinHeightTrees(self, n, edges):
        def dfs_helper(start, n):
            self.dist, self.parent = [-1]*n, [-1]*n
            self.dist[start] = 0
            dfs(start)
            return self.dist.index(max(self.dist))
        
        def dfs(start):
            for neib in Graph[start]:
                if self.dist[neib] == -1:
                    self.dist[neib] = self.dist[start] + 1
                    self.parent[neib] = start
                    dfs(neib)
                    
        Graph = defaultdict(set)
        for a,b in edges:
            Graph[a].add(b)
            Graph[b].add(a)
        
        ind = dfs_helper(0,n)
        ind2 = dfs_helper(ind, n)
        
        path = []
        while ind2 != -1:
            path.append(ind2)           #backtracking to create path
            ind2 = self.parent[ind2]
            
        Q = len(path)
        return list(set([path[Q//2], path[(Q-1)//2]]))

If you like the solution, you can upvote it on leetcode discussion section: Problem 0310