[백준] 1068. 트리/Python - Gold5

❓문제

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

성능 요약

메모리: 34936 KB, 시간: 64 ms

분류

깊이 우선 탐색, 그래프 이론, 그래프 탐색, 트리

문제 설명

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

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

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

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

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

✍🏻풀이

인접 리스트나 사전형을 활용하여 각 노드 별 자식 노드를 기록한다.

이제 지우고자 하는 노드를 삭제해야 한다. 지우려는 노드에 자식 노드가 있을 수 있기 때문에 깊이 우선 탐색으로 자식 노드를 탐색한다. 만약 지우려는 노드에 자식 노드가 있고 그 자식 노드에 또 자식 노드가 있다면 더 이상 자식 노드가 없을 때까지 탐색한다. 자식 노드가 없다면 이제 노드를 트리에서 삭제한다. 그리고 최종적으로 제거하려는 노드도 트리에서 삭제한다. 이때, 노드를 제거하며 방문한 노드들을 표시한다.

노드를 다 삭제했으니 리프 노드의 개수를 구하면 된다. 우선 방문하지 않은 노드인지 확인한다. 여기서 방문하지 않은 노드는 지워지지 않고 남아있는 노드들을 의미한다. 그리고 방문하지 않은 노드가 트리에 없다면 자식 노드가 없는 것이므로 res + 1을 한다. 이때, 트리 안에 남아있는 노드의 자식 노드에 지우고자 하는 자식 노드가 남아있다. 그러므로 이를 제외하고 남아있는 자식의 수가 0이면 res + 1을 하면 된다.

💻코드

import sys
from collections import defaultdict

input = sys.stdin.readline

N = int(input())
nodes = list(map(int, input().split()))
remove = int(input())
tree = defaultdict(list)
visited = [False] * N

def dfs(r):
    visited[r] = True
    for n in tree[r]:
        if not tree[n]:
            visited[n] = True
            tree.pop(n)
        else:
            dfs(n)
    tree.pop(r)

for v in range(N):
    if nodes[v] == -1:
        continue
    tree[nodes[v]].append(v)

dfs(remove)
res = 0

for n in range(N):
    if not visited[n]:
        if remove in tree[n] and len(tree[n])-1 == 0:
            res +=1
        if not tree[n]:
            res += 1

print(res)

📍후기

진짜 반례 생각하는 것이 매우 중요한 것 같다. 매번 테스트케이스만 통과하면 제출해서 계속 떨어지나보다..🥲

그리고 삼성 코테의 경우 pypy3인 걸로 알고 있어 pypy3로도 제출해보았다.메모리가 너무 차이나는데 어떻게 줄이면 좋을지..