[트리] 트리의 리프노드 개수 구하기

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

1068번: 트리

 

import sys
si = lambda : sys.stdin.readline().rstrip()
l = lambda : list(map(int, si().split()))

def dfs(root) : # x의 subtree에 대해 leaf[x]를 계산해주는 함수
    if not child[root] : # 아무것도 없다면 리프노드
        leaf[root] = 1

    for node in child[root] : # 자식이 존재한다면
        dfs(node) # 탐색 수행
        leaf[root] += leaf[node] # 탐색하고 나와서 리프노드의 수만큼 저장

n = int(si())
child = [[] for _ in range(n)]
leaf = [0 for _ in range(n)]
lst = l()
root = -1
for i in range(len(lst)) :
    if lst[i] == -1 :
        root = i
        continue
    child[lst[i]].append(i)
erased = int(si())

for i in range(n) : # remove할 노드를 찾아서 제거 => 탐색자체를 못하게 함.
    if erased in child[i] :
        child[i].remove(erased)

if root != erased : # 만약 타겟이 root노드가 아니라면 탐색을 수행
    dfs(root)

print(leaf[root]) # 타겟이 root라면 해당 조건을 넘어가서 0이 출력될 것