[Python] 백준 1068_ 트리

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

이 문제는 깊이 우선 탐색(dfs)을 이용하여 해결하였다.

(예제)
9
-1 0 0 2 2 4 4 6 6
4

1. 트리 그래프 만들기

먼저 루트 노드는 root변수에 따로 저장해두고, node리스트의 인덱스 값의 자식 노드 번호를 저장해두었다.

root=-1
node =[[] for _ in range(n)]
for i in range(n):
    if p[i]!=-1:
        node[p[i]].append(i)
    else: #루트일경우
        root = i  

root = 0 이다.
0의 자식노드는 1,2
1의 자식노드는 존재하지 않으므로 리프노드라고 할 수 있다.

2. 리프 노드의 개수 구하기

리프노드의 개수를 cnt변수에 저장한다.

def dfs(r): #매개변수 root
    global cnt
    if len(node[r])==0: #자식노드x, 리프노드라면
        cnt+=1
    else: #아닐 경우 이어진 자식노드 탐색 
        for i in node[r]:
            dfs(i)

3. 제거해야할 노드 제거

for i in range(n):
    if x in node[i]:
        node[i].remove(x)

제거해야할 노드인 4번 노드가 제거된 것을 확인할 수 있다.
제거할 노드를 제거하면 dfs가 실행되었을때 자연스럽게 4번노드는 탐색하지 않게 된다.

4. 리프노드 탐색 시작(dfs 실행)

cnt=0
if x!=root:
    dfs(root) #리프노드 탐색 
print(cnt)

5. 최종 완성 코드

import sys
input = sys.stdin.readline

n= int(input())
p=list(map(int, input().split()))
x =int(input())

#트리 그래프 만들기
root=-1
node =[[] for _ in range(n)]
for i in range(n):
    if p[i]!=-1:
        node[p[i]].append(i)
    else: #루트일경우
        root = i      
#print(node)

#리프노드 개수 구하기 
def dfs(r): #매개변수 root
    global cnt
    if len(node[r])==0: #자식노드x, 리프노드라면
        cnt+=1
    else: #아닐 경우 이어진 자식노드 탐색 
        for i in node[r]:
            dfs(i)

#제거해야할 노드 먼저 제거 
for i in range(n):
    if x in node[i]:
        node[i].remove(x)
#print(node)

cnt=0
if x!=root:
    dfs(root) #리프노드 탐색 
print(cnt)