[solved.ac 골드2] 2250_트리의 높이와 너비 (파이썬, DFS)

밑의 사진을 클릭하면 문제 링크로 이동합니다

제한 사항

문제

이진트리를 다음의 규칙에 따라 행과 열에 번호가 붙어있는 격자 모양의 틀 속에 그리려고 한다. 이때 다음의 규칙에 따라 그리려고 한다.

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

이와 같은 규칙에 따라 이진트리를 그릴 때 각 레벨의 너비는 그 레벨에 할당된 노드 중 가장 오른쪽에 위치한 노드의 열 번호에서 가장 왼쪽에 위치한 노드의 열 번호를 뺀 값 더하기 1로 정의한다. 트리의 레벨은 가장 위쪽에 있는 루트 노드가 1이고 아래로 1씩 증가한다.

아래 그림은 어떤 이진트리를 위의 규칙에 따라 그려 본 것이다. 첫 번째 레벨의 너비는 1, 두 번째 레벨의 너비는 13, 3번째, 4번째 레벨의 너비는 각각 18이고, 5번째 레벨의 너비는 13이며, 그리고 6번째 레벨의 너비는 12이다.

우리는 주어진 이진트리를 위의 규칙에 따라 그릴 때에 너비가 가장 넓은 레벨과 그 레벨의 너비를 계산하려고 한다. 위의 그림의 예에서 너비가 가장 넓은 레벨은 3번째와 4번째로 그 너비는 18이다. 너비가 가장 넓은 레벨이 두 개 이상 있을 때는 번호가 작은 레벨을 답으로 한다. 그러므로 이 예에 대한 답은 레벨은 3이고, 너비는 18이다.

임의의 이진트리가 입력으로 주어질 때 너비가 가장 넓은 레벨과 그 레벨의 너비를 출력하는 프로그램을 작성하시오

입력

첫째 줄에 노드의 개수를 나타내는 정수 N(1 ≤ N ≤ 10,000)이 주어진다. 다음 N개의 줄에는 각 줄마다 노드 번호와 해당 노드의 왼쪽 자식 노드와 오른쪽 자식 노드의 번호가 순서대로 주어진다. 노드들의 번호는 1부터 N까지이며, 자식이 없는 경우에는 자식 노드의 번호에 -1이 주어진다.

 

출력

첫째 줄에 너비가 가장 넓은 레벨과 그 레벨의 너비를 순서대로 출력한다. 너비가 가장 넓은 레벨이 두 개 이상 있을 때에는 번호가 작은 레벨을 출력한다.

예제

 

---

나의 풀이

위와 같은 트리가 있을 때

• InOrder는 B - A - C 순으로 탐색하는 것이고
• PreOrder는 A - B - C 순으로 탐색하는 것이고
• PostOrder는 B - C - A 순으로 탐색하는 것이다.

 

1번 열부터 19번 열까지 노드를 보면 8 - 4 - 2 - 14 - 9 - 18 - 15 - ... - 19 - 13 - 7 순으로 배치된다.

이 순서는 InOrder로 탐색한 결과이다.

그래서 주어진 트리를 InOrder로 탐색할 수 있다면 위 그림에서 열 순서(1 ~ 19)대로 탐색할 수 있는 것이다.

 

또 하나 생각해야 할 점은 예제에서는 1이 루트 노드로 주어졌지만 항상 루트 노드로 1이 주어진다고는 안 했기 때문에 루트 노드를 찾는 방법도 생각해야 한다.

이러한 것들을 고려해서 작성한 코드는 다음과 같다.

전체 코드

n = int(input())
graph = [[0, 0] for _ in range(n + 1)] # [왼쪽 자식 노드, 오른쪽 자식 노드]
row = [list() for _ in range(n + 1)] # 각 레벨의 노드들 저장
root_check = [0] * (n + 1) # 루트 노드 찾기 위한 리스트
col = 1 # row에 추가 될 노드의 열

# 그래프 연결 및 루트 노드 확인 준비
def init_graph():
    for _ in range(n):
        node, left, right = map(int, input().split())
        graph[node] = [left, right]
        root_check[node] += 1
        if left != -1: root_check[left] += 1
        if right != -1: root_check[right] += 1

# 루트 노드 찾기
def find_root():
    return root_check.index(1) # 루트 노드만 1을 가진다

# 레벨 별 노드 추가
def dfs(node, level):
    global col
    left, right = graph[node]
    if graph[node][0] != -1:
        dfs(left, level + 1)
    row[level].append(col)
    col += 1
    if graph[node][1] != -1:
        dfs(right, level + 1)

# 최대 넓이를 가진 레벨 찾아서 출력
def print_answer():
    level = 1
    max_dist = max(row[1]) - min(row[1]) + 1
    for i in range(2, n + 1):
        if not row[i]: break # 레벨이 존재하지 않으면 종료
        distance = max(row[i]) - min(row[i]) + 1
        if max_dist < distance:
            max_dist = distance
            level = i
    print(level, max_dist) # 정답 출력


init_graph()
dfs(find_root(), 1)
print_answer()

• init_graph(): 노드를 입력 받아 graph에 추가한다. root_check에도 값을 추가하는데 현재 노드와 자식 노드에 추가하고 있다. 즉, 루트 노드를 제외한 모든 노드는 현재 위치에서 추가된 1과 부모 노드가 추가한 1때문에 2를 가지게 되고, 루트 노드는 부모 노드가 없기 때문에 1이 된다.
• find_root(): root_check의 값 중 1의 위치를 반환한다. 그 위치가 루트 노드가 된다.
• dfs(): InOrder로 탐색한다. 가장 왼쪽 노드를 만날 때까지 col은 1로 유지될 것이고, row에 추가된 후에야 다음 노드를 탐색할 것이다. 재귀 함수에 대해 생각해보자.
• print_answer(): dfs()를 수행하면 row에는 각 레벨 별 몇 열에 노드가 존재하는지 저장되어 있다. 레벨 L에서의 거리는 max(L) - min(L) + 1 이다. 가장 큰 거리와 해당 레벨을 출력하며 프로그램을 종료한다.

 

채점 결과

동일한 코드에 PyPy3와 Python3로 채점한 결과이다.