Minimum Height Trees

The problem description can be found in LeetCode.

Approach

Our input Tree is given as a list of edges as input. Hence, we can build a Graph by building an Adjacency List from these edges.

Another thing to consider is that this graph doesn't contain cycles. Moreover, there is just a single path between any two nodes of the Graph itself.

So, if we think about our input Tree as a skewed one, the internal nodes would have the minimum height if rooted, whereas the leaf nodes would have the maximum one.

Based on this assumption, we can design an algorithm to look for such internal nodes in the skewed tree, which would give us the minimum height.

Solution

from collections import deque
from typing import Deque, Dict, List


class Solution:
    def findMinHeightTrees(self, n: int, edges: List[List[int]]) -> List[int]:
        assert n > 0
        if n == 1:
            return [0]
        if n == 2:
            return [0, 1]

        graph = self.get_adjacency_list(n, edges)

        return self.topological_sort(n, graph)

    def get_adjacency_list(
        self, n: int, edges: List[List[int]]
    ) -> Dict[int, List[int]]:
        graph: Dict[int, List[int]] = {i: [] for i in range(n)}

        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)

        return graph

    def topological_sort(self, n: int, graph: Dict[int, List[int]]) -> List[int]:
        queue: Deque[int] = deque([])

        # add leaf nodes
        for node, neighbors in graph.items():
            if len(neighbors) == 1:
                queue.append(node)

        not_processed_nodes = n
        while not_processed_nodes > 2:
            not_processed_nodes -= len(queue)

            for _ in range(len(queue)):
                node = queue.popleft()
                neighbor = graph[node].pop()
                graph[neighbor].remove(node)

                # the node is now a leaf one
                if len(graph[neighbor]) == 1:
                    queue.append(neighbor)

        return list(queue)

Complexity

• Time: $O(n)$
• Space: $O(n)$

Where $n$ is the number of nodes in the graph.