[BOJ/백준] 1068번 트리 (Python 파이썬)

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

1068번: 트리

 

문제

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

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

예를 들어, 다음과 같은 트리가 있다고 하자.

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

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

입력

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

출력

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

예제 입력 1 

5
-1 0 0 1 1
2

예제 출력 1 

2

예제 입력 2 

5
-1 0 0 1 1
1

예제 출력 2 

1

예제 입력 3 

5
-1 0 0 1 1
0

예제 출력 3 

0

예제 입력 4 

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

예제 출력 4 

2

 

---

풀이

문제를 살펴보면 입력으로 각 노드들의 부모 정보가 주어지고, 제거할 노드의 번호가 주어집니다.

따라서 우선 (1) 주어진 부모들의 정보로 tree를 그리고, (2) 노드를 삭제하는 연산을 구현해야 합니다.

우선 tree를 그리는 과정부터 살펴보면 아래와 같습니다.

N = int(input())
parent_info = list(map(int, input().split()))

# 각 부모당 자식의 정보들 -> 자식의 개수가 0인 node가 leaf node
tree = [[] for _ in range(N)]

for i in range(N):
    if parent_info[i] != -1:
        tree[parent_info[i]].append(i)  # 자식 추가

이러면, 각 인덱스에는 자신들의 자식의 번호들이 담기게 될 것입니다.

이제 노드를 삭제하는 연산을 구현합시다. 이를 위해서 각 노드들의 생존 여부를 저장하는 리스트를 하나 선언하였습니다.

alive = [True for _ in range(N)]

만약 삭제가 일어난다면, alive에서 해당 인덱스의 값을 False로 바꿔줄 것입니다.

이제 삭제되는 노드에 대해서 경우의 수를 나눠보겠습니다.

1. internal 노드인 경우
   => 이 경우, 해당 노드의 모든 자식들까지 함께 삭제해야 합니다. 따라서 그래프 탐색을 통해 모든 자식들을 찾아서 alive을 False로 만들어 줍니다. 저는 BFS을 사용했습니다.
2. leaf 노드인 경우
   => 이 경우, 자신의 부모에서 자신이 자식이라는 정보를 지운 후 삭제되어야 합니다. 따라서 우선 부모가 누구인지 찾은 후, 해당 부모에서 자신을 지운 후 alive을 False로 만들어 줍니다.
3. root 노드인 경우
   => 이 경우, 모든 노드들이 삭제되므로 답은 0으로 고정됩니다.

우선 루트 노드을 따로 지정하기 위해 위의 코드들을 아래와 같이 수정하겠습니다.

N = int(input())
parent_info = list(map(int, input().split()))

# 각 부모당 자식의 정보들 -> 자식의 개수가 0인 node가 leaf node
tree = [[] for _ in range(N)]
# 생존 여부
alive = [True for _ in range(N)]
# root 노드 인덱스
root = 0


for i in range(N):
    if parent_info[i] != -1:
        tree[parent_info[i]].append(i)  # 자식 추가
    else:
        root = i

 

그리고 위 세가지 경우의 수를 처리하면 아래와 같습니다.

def solve(removed_node):
    cnt = 0
    alive[removed_node] = False

    # root 노드를 삭제하는 경우
    if removed_node == root:
        return 0
    # 리프노드를 삭제하는 경우
    elif len(tree[removed_node]) == 0:
        # 삭제하려는 노드의 부모 인덱스 구하기
        parent_index = 0
        for i, parent in enumerate(tree):
            for child in parent:
                if child == removed_node:
                    parent_index = i
        # 삭제하려는 노드을 부모에서 제거
        tree[parent_index].remove(removed_node)
    # internal 노드를 삭제하는 경우
    else:
        for child in tree[removed_node]:
            BFS(child)

    for i in range(N):
        if alive[i] and len(tree[i]) == 0:
            cnt += 1
    return cnt

위 3가지 경우을 모두 거친 후, (1) 아직 살아있으면서 [alive is True] (2) 자식들이 없는 경우 [len(tree) is 0] 가 아직 살아있는 리프 노드이므로 이를 세어줍니다.

따라서, BFS까지 추가한 최종적인 코드는 아래와 같습니다.

Code

import sys
from collections import deque
input = sys.stdin.readline

N = int(input())
parent_info = list(map(int, input().split()))
Removed = int(input())

# 각 부모당 자식의 정보들 -> 자식의 개수가 0인 node가 leaf node
tree = [[] for _ in range(N)]
# 생존 여부
alive = [True for _ in range(N)]
# root 노드 인덱스
root = 0


for i in range(N):
    if parent_info[i] != -1:
        tree[parent_info[i]].append(i)  # 자식 추가
    else:
        root = i


def BFS(p):
    queue = deque([p])
    # 방문 여부
    visited = [False for _ in range(N)]
    visited[p] = True
    alive[p] = False

    while queue:
        p = queue.popleft()

        for c in tree[p]:
            if not visited[c]:
                queue.append(c)
                visited[c] = True
                alive[c] = False


def solve(removed_node):
    cnt = 0
    alive[removed_node] = False

    # root 노드를 삭제하는 경우
    if removed_node == root:
        return 0
    # 리프노드를 삭제하는 경우
    elif len(tree[removed_node]) == 0:
        # 삭제하려는 노드의 부모 인덱스 구하기
        parent_index = 0
        for i, parent in enumerate(tree):
            for child in parent:
                if child == removed_node:
                    parent_index = i
        # 삭제하려는 노드을 부모에서 제거
        tree[parent_index].remove(removed_node)
    # internal 노드를 삭제하는 경우
    else:
        for child in tree[removed_node]:
            BFS(child)

    for i in range(N):
        if alive[i] and len(tree[i]) == 0:
            cnt += 1
    return cnt


print(solve(Removed))