Minimum Height Trees

Problem: Minimum Height Trees

Since the graph is a tree, we can keep removing leaf nodes in the same level until there's only one or two nodes left. What's rest will definitely be our answer since they are evenly close to all nodes.

Code in Python:

from collections import defaultdict

class Solution(object):
    def findMinHeightTrees(self, n, edges):
        """
        :type n: int
        :type edges: List[List[int]]
        :rtype: List[int]
        """
        adj = defaultdict(set)
        for u, v in edges:
            adj[u].add(v)
            adj[v].add(u)
        ans = set(range(n))
        while len(ans) > 2:
            leaves = set(i for i in ans if len(adj[i]) == 1)
            ans -= leaves
            for u in leaves:
                for v in adj[u]:
                    adj[v].remove(u)
        return list(ans)