[알고리즘] Leetcode #310 최소 높이 트리 (Python)

---

↓↓↓ 아래는 내 리트코드 계정 ↓↓↓

leetcode.com/Jiwon_Lee/

Jiwon Lee - LeetCode Profile

 

---

문제 설명

A tree is an undirected graph in which any two vertices are connected by exactly one path. In other words, any connected graph without simple cycles is a tree.

Given a tree of n nodes labelled from 0 to n - 1, and an array of n - 1 edges where edges[i] = [ai, bi] indicates that there is an undirected edge between the two nodes ai and bi in the tree, you can choose any node of the tree as the root. When you select a node x as the root, the result tree has height h. Among all possible rooted trees, those with minimum height (i.e. min(h))  are called minimum height trees (MHTs).

Return a list of all MHTs' root labels. You can return the answer in any order.

The height of a rooted tree is the number of edges on the longest downward path between the root and a leaf.

 

입출력 예

Example 1:

Input: n = 4, edges = [[1,0],[1,2],[1,3]] Output: [1] Explanation: As shown, the height of the tree is 1 when the root is the node with label 1 which is the only MHT.

Example 2:

Input: n = 6, edges = [[3,0],[3,1],[3,2],[3,4],[5,4]] Output: [3,4]

Example 3: Input: n = 1, edges = [] Output: [0]

Example 4: Input: n = 2, edges = [[0,1]] Output: [0,1]

 

Constraints

• 1 <= n <= 2 * 104
• edges.length == n - 1
• 0 <= ai, bi < n
• ai != bi
• All the pairs (ai, bi) are distinct.
• The given input is guaranteed to be a tree and there will be no repeated edges.

 

Code

# 가운데 노드 찾기
def findMinHeightTrees(n = 4, edges = [[1, 0], [1, 2], [1, 3]]):
    if n <= 1:
        return [0]
    
    # 양방향 그래프 구성
    graph = collections.defaultdict(list)
    for i, j in edges:
        graph[i].append(j)
        graph[j].append(i)
    
    # 첫번째 리프 노드 추가
    leaves = []
    for i in range(n + 1):
        if len(graph[i]) == 1:
            leaves.append(i) # 리프 노드 할당

    while n > 2:
        n -= len(leaves)
        new_leaves = []
        for leaf in leaves:
            neighbor = graph[leaf].pop()
            graph[neighbor].remove(leaf) # 그래프에 제거 작업 반영

            if len(graph[neighbor]) == 1:
                new_leaves.append(neighbor) # 리프 노드 할당

        leaves = new_leaves
        
    return leaves

입력이 무방향 트리로 주어지기 때문에 루트-리프 관계를 한 번씩 더 설정해준다. 이 문제에서는 트리의 최소 높이를 찾아주어야 하므로, 남은 노드의 개수가 2개가 될 때까지 리프노드부터 하나씩 삭제해가며 가운데 위치한 노드를 찾아줘야 한다. 이를 위해, 위에서 만들어준 graph의 길이가 1인 노드를 leaves 리스트에 할당하여 기존 두 번씩 설정한 그래프 연결관계에서 제거해준다. while문에서 이 과정을 계속 반복하다가 n > 2 조건을 깨는 순간 반환하면 최소 높이를 제공하는 노드를 찾아낼 수 있다.