[Python/파이썬] 백준 22856번 트리 순회

22856번: 트리 순회

 

문제

노드가 N개인 이진 트리가 있다. 트리를 중위 순회와 유사하게 순회하려고 한다. 이를 유사 중위 순회라고 하자.

순회의 시작은 트리의 루트이고 순회의 끝은 중위 순회할 때 마지막 노드이다. 이때 루트 노드는 항상 1번 노드이다.

유사 중위 순회는 루트 노드에서 시작하며, 다음과 같이 진행된다.

1. 현재 위치한 노드의 왼쪽 자식 노드가 존재하고 아직 방문하지 않았다면, 왼쪽 자식 노드로 이동한다.
2. 그렇지 않고 현재 위치한 노드의 오른쪽 자식 노드가 존재하고 아직 방문하지 않았다면, 오른쪽 자식 노드로 이동한다.
3. 그렇지 않고 현재 노드가 유사 중위 순회의 끝이라면, 유사 중위 순회를 종료한다.
4. 그렇지 않고 부모 노드가 존재한다면, 부모 노드로 이동한다.
5. 유사 중위 순회를 종료할 때까지 1 ~ 4를 반복한다.

위 그림에 있는 트리에서 중위 순회를 한다면 4→2→5→1→6→3→7 순으로 순회를 한다.

따라서, 유사 중위 순회의 끝은 노드 7이 된다.

유사 중위 순회는 위 그림과 같이 루트인 노드 1$1$에서 시작하여 노드 7$7$에서 끝나고 1→2→4→2→5→2→1→3→6→3→7 이와 같은 순서로 순회를 진행한다. 유사 중위 순회를 진행하면서 총 10번 이동하였다.

여기서 이동이라는 것은 하나의 노드에서 다른 노드로 한번 움직이는 것을 의미한다. 예를 들면, 노드 1에서 노드 2로 가는 것을 한번 이동하였다고 한다.

유사 중위 순회를 하면서 이동한 횟수를 구하려고 한다.

입력

첫 번째 줄에 트리를 구성하는 노드의 개수 N이 주어진다.

두 번째 줄부터 N + 1번째 줄까지 현재 노드 a, 현재 노드의 왼쪽 자식 노드 b, 현재 노드의 오른쪽 자식 노드 c가 공백으로 구분되어 주어진다. 만약 자식 노드의 번호가 -1인 경우 자식 노드가 없다는 것을 의미한다.

출력

유사 중위 순회를 하면서 이동한 총 횟수를 출력한다.

제한

•  1 ≤ N ≤ 100,000
•  1 ≤ a, b ≤ 1

 

코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(100000)

n = int(input())
tree = [[] for _ in range(n+1)]
visited = [False] * (n+1)
ans = 0

def inorder(node):
    global ans
    # 왼쪽 노드
    if tree[node][0] != -1 and not visited[tree[node][0]]:
        ans += 1
        inorder(tree[node][0])
    # 현재 노드 방문 처리
    visited[node] = True
    # 현재 노드가 유사 중위 순회의 끝인 경우
    if node == lastNode:
        print(ans)
        return
    # 오른쪽 노드
    if tree[node][1] != -1 and not visited[tree[node][1]]:
        ans += 1
        inorder(tree[node][1])
    ans += 1
        
for _ in range(n):
    a, b, c = map(int, input().split())
    tree[a] = [b, c]

# 가장 마지막으로 방문하는 노드 찾기
now = 1
lastNode = 1
while True:
    if tree[now][1] == -1:
        lastNode = now
        break
    now = tree[now][1]
    
inorder(1)

 

가장 마지막으로 방문하는 노드는 트리의 가장 오른쪽 노드일 것이다. 그러므로 오른쪽 자식이 없는 노드가 나올 때까지 오른쪽으로 계속 탐색하면 마지막으로 방문하는 노드(lastNode)를 찾을 수 있다.