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

1068번: 트리

 

코드

import sys

n = int(input())
tree = [[] for _ in range(n)]
p = list(map(int,sys.stdin.readline().split()))
for i in range(n):
    tree[i] = p[i]
remove = int(input())

def delete(remove):
    tree[remove] = -2
    for i in range(n):
        if tree[i] == remove:
            tree[i] = -2
            delete(i)
delete(remove)
cnt = 0
for i in range(n):
    if tree[i] != -2:
        err = 0
        for j in tree:
            if j == i:
                err = 1
                break
        if err == 0:
            cnt += 1
print(cnt)

 

설명

tree 리스트에 노드 번호에 따라 각자의 부모 노드를 저장해주었다. 

delete 함수에서는 먼저 지워지는 노드(remove)의 부모 노드 값을 -2로 바꿔줬다. remove 노드 밑에 있는 모드 노드들도 같이 제거되기 때문에 dfs를 통해 아래에 이어진 모든 노드들의 값도 -2로 바꿔 삭제해줬다.

tree 리스트를 보고 만약 값이 -2가 아니고 자신을 부모 노드로 갖는 노드도 존재하지 않는다면 cnt를 1씩 증가시켰다.