[백준] 트리의 높이와 너비 (파이썬, python)

🔗 트리의 높이와 너비

문제

1. 이진 트리에서 같은 레벨에 있는 노드는 같은 행에 위치한다.
2. 한 열에는 한 노드만 존재한다.
3. 임의의 노드의 왼쪽 subtree에 있는 노드들은 해당 노드보다 왼쪽 열에 위치하고, 오른쪽 서브트리에 있는 노드들은 해당 노드보다 오른쪽 열에 위치한다.
4. 노드가 배치된 가장 왼쪽 열과 오른쪽 열 사이엔 아무 노드도 없이 비어있는 열은 없다.

이미지 출처: https://www.acmicpc.net/problem/2250

레벨의 너비는 해당 레벨의 노드 중 가장 왼쪽에 있는 노드와 가장 오른쪽에 있는 노드의 차 + 1 이다.

임의의 이진트리가 입력으로 주어질 때 너비가 가장 넓은 레벨과 그 레벨의 너비를 출력해라.

⏱ 노드의 개수 N (1≤ N ≤ 10000)

🔑 풀이

해당 문제는 알고리즘 자체는 어렵지 않았지만, 저장해야 할 값들이 많아서 헷갈렸다.

각 레벨(깊이)의 너비와 해당 레벨 값을 구하기 위해서

1. 입력에 따라 트리 자료 구조를 만들고
2. 해당 트리 탐색하며 각 레벨의 가장 왼쪽 idx, 가장 오른쪽 idx를 갱신 해야한다.

1. 트리 구현, 루트 노드 값 찾기

dict를 활용하여 입력을 트리로 만들었다.

    tree = {1 : [2,3]} # root : [left_node, right_node] 

• 트리를 탐색할 때 root 부터 시작 해야하는데 root가 항상 1이 아닐 수 있기 때문에 각 노드의 부모 노드를 저장하는 배열을 새로 만들어 부모 노드 값을 저장해야한다.
• N개의 노드가 주어졌을 때, 노드의 값은 1~ N 까지 매겨지기 때문에 parentNode 라는 배열을 만들어 해당 idx에 부모노드 idx를 저장했다.

2. 중위 순회

노드를 살펴보면, 8→ 4 → 2 → 14 → 9 → 18 → .. 순으로 idx가 1부터 N까지 매겨진다.

• 왼쪽 노드 → 루트 노드 → 오른쪽 노드
  이 순서는 순으로 탐색하는 중위순회 방법이라는 걸 알 수 있다.

가장 왼쪽에 있는 노드 부터 가장 오른쪽에 있는 노드까지 오른쪽으로 하나씩 순차적으로 탐색하며 아래 값을 구한다.

• 현재 레벨의 왼쪽, 오른쪽 노드의 레벨을 저장
• 현재 레벨의 가장 왼쪽, 오른쪽 idx 갱신

    def in_order(key):
        global cur
        if tree[key][0] != -1: # 왼쪽 노드가 있을 때
            level[tree[key][0]] = level[key] + 1 # 왼쪽 노드의 레벨은 현재 레벨+1
            in_order(tree[key][0])
    
        cur += 1 # 노드 방문하면서 1 부터 N까지 idx 증가
        
        if level_min[level[key]] > cur: 
            level_min[level[key]] = cur # 현재 레벨의 가장 왼쪽 노드 idx 갱신
        if level_max[level[key]] < cur:
            level_max[level[key]] = cur
    
        if tree[key][1] != -1:
            level[tree[key][1]] = level[key] + 1
            in_order(tree[key][1])
    
        return

• level [] 에는 각 노드의 깊이를 저장한다. ( idx는 노드 번호, value는 깊이 )
• level_min에는 각 레벨의 최소 idx를 저장한다. (idx는 레벨, value는 cur로 갱신)

3. level_min, level_max 순회하며 최대 차이 레벨 구하기

코드

    n = int(input())  
    parentNode = [-1]*(n+1)
    tree = {}
    root = 0
    
    for _ in range(n):  # 1) 트리 형태 만들기
        v, left, right = map(int, input().split(' '))
        tree[v] = [left, right]
        if left != -1:
            parentNode[left] = v
        if right != -1:
            parentNode[right] = v
    
    for i in range(1, n+1):  # 2) 루트 노드 찾기
        if parentNode[i] == -1:
            root = i
    
    level_min = [n]*(n+1)  # 각 레벨의 최소, 최대 idx 저장
    level_max = [0]*(n+1)
    level = [0]*(n+1)
    level[root] = 1
    cur = 0
    
    
    def in_order(key): # 3) 중위 순회
        global cur
        if tree[key][0] != -1:
            level[tree[key][0]] = level[key] + 1
            in_order(tree[key][0])
    
        cur += 1
        
        if level_min[level[key]] > cur:
            level_min[level[key]] = cur
        if level_max[level[key]] < cur:
            level_max[level[key]] = cur
    
        if tree[key][1] != -1:
            level[tree[key][1]] = level[key] + 1
            in_order(tree[key][1])
    
        return
    
    
    in_order(root)
    
    result = 0
    result_idx = 0
    
    for i in range(1, n+1): # 4) 차이 찾기
        diff = level_max[i]-level_min[i]+1
        if result < diff:
            result = diff
            result_idx = i
    
    print(result_idx, result)