백준 1068번 : 트리 in Python

1068번: 트리

 

www.acmicpc.net

이 문제는 트리에서 노드 하나를 지웠을 때 남은 트리에서 리프 노드의 개수를 구하는 문제이다.

이 문제는 트리를 각 노드의 부모 노드를 표현하는 배열로 나타내기 때문에 노드 하나를 단순히 제거하기는 쉽지 않다. 노드 하나를 제거할 때 그 노드의 자식 노드 역시 제거되기 때문에, 이를 부모 노드를 표현하는 배열에 반영해야 한다.

노드 하나를 제거하는 함수를 구현해보자. 우선 각 노드의 부모 노드를 표현하는 배열을 parents라 하고, 노드 u를 제거하는 경우 이를 parents 배열에는 parents[u] = -2로 반영한다고 하자. 그러면 노드 v를 제거하는 함수에서는 우선 parents[v]를 -2로 하고, 모든 노드를 순회하면서 만약 부모 노드가 v인 노드 i를 찾으면 이 함수를 i를 인수로 해서 재귀적으로 호출하면 될 것이다. 그러면 노드 v를 루트로 하는 서브트리의 모든 노드를 제거할 수 있다.

이제 모든 노드에 대해 순회하면서 리프 노드를 세면 된다. 리프 노드는 제거된 노드가 아니면서, 다른 노드의 부모 노드가 아니면 된다. 즉, 노드 i에 대해 parents[i]가 -2가 아니고 i가 parents의 원소가 아니면 된다.

이를 코드로 구현하면 다음과 같다.

N = int(input())
parents = list(map(int, input().split()))
erase = int(input())
# 노드를 제거하는 함수
def erasing(u):
    parents[u] = -2
    for i in range(N):
        if u == parents[i]:
            erasing(i)
# 리프 노드의 개수
cnt = 0
erasing(erase)
# 제거된 노드가 아니면서 어떤 노드의 부모 노드가 아닌 경우가 트리의 리프 노드이다.
for i in range(N):
    if parents[i] != -2 and i not in parents:
        cnt += 1
print(cnt)