문제

---

트리에서 리프 노드란, 자식의 개수가 0인 노드를 말한다.

트리가 주어졌을 때, 노드 하나를 지울 것이다. 그 때, 남은 트리에서 리프 노드의 개수를 구하는 프로그램을 작성하시오. 노드를 지우면 그 노드와 노드의 모든 자손이 트리에서 제거된다.

이제 리프 노드의 개수는 1개이다.

입력

---

첫째 줄에 트리의 노드의 개수 N이 주어진다. N은 50보다 작거나 같은 자연수이다. 둘째 줄에는 0번 노드부터 N-1번 노드까지, 각 노드의 부모가 주어진다. 만약 부모가 없다면 (루트) -1이 주어진다. 셋째 줄에는 지울 노드의 번호가 주어진다.

출력

---

첫째 줄에 입력으로 주어진 트리에서 입력으로 주어진 노드를 지웠을 때, 리프 노드의 개수를 출력한다.

예제 입출력

예제 입력 1 
5
-1 0 0 1 1
2

예제 출력 1 
2

Code

from collections import defaultdict
graph = defaultdict(list)
N = int(input())
L = list(map(int,input().split()))
num = int(input())
leafnode = 0
for i in range(N) :
    if L[i] == -1 :
        leafnode = i
    else : 
        graph[L[i]].append(i)
        graph[i].append(L[i])
visited = [0] * N
if num == leafnode :
    print(0)
    exit()
#print(num)
for i in range(N) :
    try :
        graph[i].remove(num)
    except : pass
leaf = visited.copy()
def DFS(grpah, node) : 
    visited[node] = 1
    if node == num : return 0
    for v in graph[node] :
        if not visited[v] :
            if len(graph[v]) <= 1 :
                leaf[v] = 1
            DFS(graph, v)
#print(graph)
DFS(graph, leafnode)
#print(leaf)
ans = sum(leaf)
if ans == 0 :
    print(1)
else : print(ans)

해설

내가 짠 코드 기준으로 루트 노드를 제외한 모든 딕셔너리 value의 length가 1인 key는 리프 노드가 된다. 따라서 루트 노드가 터미널 노드가 될 수 있다던가..하는 예외처리만 잘 해주면 올바른 값이 나온다.