문제 링크

<aside> 💡 깊이 우선 탐색/그래프 이론/그래프 탐색/트리

</aside>

Memory 31256KB Time 44ms Code Length 372B

import sys
input = sys.stdin.readline

def dfs(n, arr):
    arr[n] = -2
    for i in range(len(arr)):
        if n == arr[i]:
            dfs(i, arr)

n = int(input())
parents = list(map(int, input().split()))
remove = int(input())

dfs(remove, parents)

count = 0
for i in range(len(parents)):
    if parents[i] != -2 and i not in parents:
        count += 1
print(count)

이 코드는 트리의 노드를 제거하고, 그 결과로 생성된 새로운 트리의 개수를 세는 것이다.

먼저, sys.stdin.readline을 이용해 입력을 받는다. 이후에는 입력받은 노드들의 부모 노드 정보를 parents 리스트에 저장한다.

remove 변수에는 제거할 노드의 번호를 저장하고, dfs 함수를 호출하여 제거할 노드와 그 하위 노드들을 모두 -2로 표시한다.

마지막으로, parents 리스트를 순회하면서 -2가 아닌 노드 중에서 자신이 부모 노드 리스트에 없는 노드, 즉 루트 노드인 경우를 세어서 출력한다. 이것이 새로 생성된 트리의 개수가 된다.