[BOJ][🟡5][백준#01068] 트리

문제 출처

BAEKJOON Online Judge #1068

문제

트리에서 리프 노드란, 자식의 개수가 0인 노드를 말한다. 트리가 주어졌을 때, 노드 하나를 지울 것이다. 그 때, 남은 트리에서 리프 노드의 개수를 구하는 프로그램을 작성하시오. 노드를 지우면 그 노드와 노드의 모든 자손이 트리에서 제거된다. 예를 들어, 다음과 같은 트리가 있다고 하자.

현재 리프 노드의 개수는 3개이다. (초록색 색칠된 노드) 이때, 1번을 지우면, 다음과 같이 변한다. 검정색으로 색칠된 노드가 트리에서 제거된 노드이다.

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

입력

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

출력

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

예제

예제 1

입력

5
-1 0 0 1 1
2

출력

2

예제 2

입력

5
-1 0 0 1 1
1

출력

1

예제 3

입력

5
-1 0 0 1 1
0

출력

0

예제 4

입력

9
-1 0 0 2 2 4 4 6 6
4

출력

2

My Sol

def main():
    def find_root():
        for n in range(N):
            if arr[n]==-1: return n

    def find_children(p):
        if p == remove_node: return 0
        ssum = 0
        found_remove_node = False
        for c in children[p]:
            if c == remove_node:
                found_remove_node = True
                continue
            ssum += find_children(c) if children[c] else 1

        if found_remove_node and ssum==0: return 1
        return ssum

    N = int(input())
    arr = list(map(int, input().split()))
    remove_node = int(input())

    children = [[] for _ in range(N+1)]
    for n in range(N):
        children[arr[n]].append(n)

    root = find_root()
    return find_children(root)

print(main())

트리의 부모자식 노드 구조를 고려해 DFS로 풀 수 있는 문제였다.

1. 입력을 받는다. 노드들의 자식들을 arr로, 지우는 노드를 remove_node로 명명했다.
2. 자식 노드를 의미하는 children 2차원 배열을 두고, arr의 각 노드들에서 부모를 인덱스로 하는 list에 자식노드의 번호를 추가해 넣었다. 이렇게 되면, 리프노드는 빈 리스트를 가지게 된다.
3. find_root 함수를 통해 루트 함수를 찾는데, 이는 arr에서 -1인 노드를 찾는 것이다. 이를 root라고 명명했다.
4. find_children 함수에 root부터 시작하는 DFS를 시작한다. find_children 함수는 parameter로 전달받은 번호의 노드로부터 자식 노드를 순회하며 조회하는데, 기준이 되는 부모 노드가 제거되는 노드라면 리프노드를 따질 수 없기 때문에 0을 반환한다.
5. 자식 노드를 매번 마찬가지로 find_children 함수를 따라가며 빈 리스트가 나오면 1을 추가해 반환하는 것이다.
6. 그런데 이때 1자 흐름이라면 remove_node를 맞이했을 때 바로 0이 되는 것이 아닌, 그 이전의 부모노드가 리프노드가 되므로 found_remove_node 인자를 확인하여 remove_node만을 자식 노드로 가지는 부모 노드가 있는 경우 1을 반환하도록 하였다.

결과

맞았습니다!!