[BOJ] 1068 트리 | Python

문제

백준 1068번 트리 | 골드5

https://www.acmicpc.net/problem/1068

1068번: 트리

• 알고리즘 분류: 그래프 이론 | 그래프 탐색 | 트리 | 깊이 우선 탐색

 

 

 

풀이 과정

파이썬 풀이

import sys
input = sys.stdin.readline

def dfs(x):
    lst[x] = -2
    for i in child[x]:  # 자식 중 제거되지 않은 노드에 대해 재귀
        if lst[i] != -2:
            dfs(i)

n = int(input())
lst = list(map(int, input().split()))
child = [[] for _ in range(n)]
for i in range(n):
    if lst[i] == -1:    # 파이썬에서는 -1이면 마지막 인덱스에 추가됨 
        continue
    child[lst[i]].append(i)

remove = int(input())
dfs(remove)
cnt = 0

for i in range(n):
    flag = True
    if lst[i] == -2:
        flag = False
        continue

    for j in child[i]:      # 자식 중 하나라도 -2가 아닌 게 있다면
        if lst[j] != -2:    # 자식이 있는 것이므로 flag를 False로
            flag = False
            break
    if flag:                # flag가 True이면 자식이 없음
        cnt += 1

print(cnt)

• DFS 사용
• lst에 부모 노드가 저장되므로, 이를 순회하며 child라는 배열을 통해 각 노드별 자식을 저장
  이때 lst의 값이 -1이라면 루트 노드로, 해당 노드를 자식으로 추가하면 안되는데 파이썬의 경우 -1을 마지막 인덱스로 하여 값을 추가하기 때문에 이 경우 continue를 사용해 저장되지 않도록 하였다.
• 처음에 0번 노드가 루트 노드라 생각해 1부터 n-1까지 child에 추가하도록 하였는데, 항상 0번 노드가 루트 노드인 것은 아니기 때문에 인덱스로 판단하는 것이 아니라 lst[i]가 -1인지를 기준으로 판단해야 한다.
• dfs 함수에서는 lst에서 해당 인덱스의 값을 -2로 하여 지워진 노드임을 표시하였고, 재귀를 통해 자식 노드들 역시 표시
• lst의 값이 -2가 아니고(지워진 노드가 아님), 지워지지 않은 자식 노드가 하나도 없는 경우 cnt를 추가하여 출력