[백준] Python - 트리

1068번: 트리

 

트리의 리프 노드의 개수를 확인하는 문제다.

내 개인 서버에서는 잘 돌아가는데 백준 테스트 케이스에서 어떤 경우가 있는지 몰라서 예제만 확인했더니...

에러를 해결하지 못했다.

 

그래서 어떤 알고리즘으로 찾아 힌트를 봤더니 dfs를 쓴다 했다.

깊이 우선 탐색을 재귀로 짜서 리프노드에 도달하면 갯수 +1 하고 리턴, 삭제노드의 부모가 자식이 1개면 갯수 +1 이런식으로 구현했다.

from sys import stdin

n = int(stdin.readline())
parent = list(map(int, stdin.readline().split()))
d = int(stdin.readline())
tree = [[] for _ in range(n)]
leafCount = 0

for i in range(0, n):
    if parent[i] != -1:
        tree[parent[i]].append(i)
    else:
        root = i

def dfs(node):
    global leafCount, d

    if len(tree[node]) == 0 : 
        leafCount += 1
        return
    else:
        for i in tree[node]:
            if i == d and len(tree[node]) == 1: 
                leafCount += 1
            elif i == d : 
                continue
            else: 
                dfs(i)

if d == root: 
    print(0)
else:
    dfs(root)
    print(leafCount)