[백준/파이썬]1068 트리

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

1068번: 트리

문제

백준 1068번 트리

이진트리가 주어졌을 때, 리프 노드라는 것의 개수를 구하는 문제다.

 

나는 이 문제를 풀 때 입력으로 주어지는 부모 노드를 이용해 다음과 같이 인접 리스트를 만들어 준 후에 삭제할 노드를 삭제하고 리프 노드 카운팅을 해줬다.

path = [[] for _ in range(n)]

for i in range(n) :
    if node[i] == -1 :
        root_node = i
    else :
        path[node[i]].append(i)

path[del_node] = []

 

 

그리고 이 문제를 풀고 처음 제출했을 때 78%에서 틀렸습니다가 나왔다.

 

알고 보니 트리가 일자로 쭉 이어졌을 경우 (0 -> 1 -> 2 -> 3)를 처리해주지 못해서 틀렸다고 나왔었다.

 

정답

def dfs(node_idx) :
    global count
    if len(path[node_idx]) == 0 and node_idx != root_node:
        count+=1
        return
    elif len(path[node_idx]) == 1 and path[node_idx][0] == del_node :
        count += 1
        return
    for child in path[node_idx] :
        if child != del_node :
            dfs(child)

n = int(input())
node = list(map(int, input().split()))
del_node = int(input())

path = [[] for _ in range(n)]

for i in range(n) :
    if node[i] == -1 :
        root_node = i
    else :
        path[node[i]].append(i)

path[del_node] = []

count = 0
dfs(root_node)

print(count)