[Python | 파이썬] 백준 1068 트리

[Python | 파이썬] 백준 1068 트리 

 

0번부터 n-1번 노드에 대해서 루트노드이면 -1, 그렇지 않으면 부모노드의 번호가 입력으로 한 줄에 주어진다.

-1이면 root에 저장해주고, -1이 아니라면 부모노드가 있는 노드이다.

이 문제에서 트리의 간선은 양방향 간선이므로 s[i][arr[i]] = True, s[arr[i]][i] = True로 초기화해준다.

그리고 지울 노드의 번호와 연결된 간선들은 모두 False로 갱신해준다.

이후 dfs를 이용해 모든 노드를 탐색하는데 방문하는 노드의 visit배열을 True로 설정해준다.

그리고 flag 불리언 값을 False로 초기화하고, 현재 노드에 자식노드가 한개라도 있고 그 자식노드를 방문하지 않았으면 flag를 True로 바꿔주고, 재귀적으로 탐색한다. 만약 flag가 False라면 count값을 1증가 시킨다.

재귀탐색을 종료한 이후에 target이 root노드였다면 전부 지워지니까 0을 출력하고, 그렇지않다면 count를 출력한다.

def dfs(root):
    global cnt
    flag = False
    visit[root] = True
    for i in range(n):
        if s[root][i] and not visit[i]:
            flag = True
            dfs(i)
    if not flag:
        cnt += 1

n = int(input())
arr = list(map(int, input().split()))
target = int(input())
s = [[0] * n for _ in range(n)]
visit = [False] * n
length = len(arr)
cnt = 0
for i in range(length):
    if arr[i] != -1:
        s[i][arr[i]] = True
        s[arr[i]][i] = True
    else:
        root = i

for i in range(n):
    s[i][target] = False
    s[target][i] = False

dfs(root)
if target == root:
    print(0)
else:
    print(cnt)

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