백준[ 1068 | Python ] 트리

문제

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

풀이

1068에서는 먼저 배열이 주어진다.
이를 DFS로 탐색하기에 용이하게 하기 위해서 서로의 연결정보를 알 수 있는 그래프의 형태로 만들어준다.

이후로 DFS를 통해 탐색하는데 노드의 개수를 세어주는 것이 아닌 리프의 개수를 세어줘야하기 때문에 어떤 노드 s를 방문했을 때 그 노드가 더 다른 연결정보를 가지고 있지 않다면 리프로 취급하여 수를 세어준다.

또한 이전에 받은 제거할 노드의 정보를 사용해 현재 탐색하는 노드가 삭제할 노드일 경우(=target) 그 노드가 삭제되었을 때 target노드는 리프가 되기 때문에 리프의 개수를 올려준다.

위의 조건들이 충족되지 않은경우 DFS를 통해서 탐색을 이어나간다.

이 때 DFS를 n부터 시작하는 이유는 부모가 없는 노드 즉, 루트 노드를 그래프를 만들어줄 때 n에 저장했기 때문이다.

import sys
sys.setrecursionlimit(10**6)
n = int(input())
lst = list(map(int,input().split()))
target = int(input())
graph = [[] for _ in range(n+1)]

for i in range(n):
    if lst[i] == -1:
        graph[n].append(i)
    else:
        graph[lst[i]].append(i)

def dfs(s):
    global leaves
    if not graph[s]:
        leaves+=1     
    for i in range(len(graph[s])):
        if graph[s][i] == target:
            if len(graph[s]) == 1:
                leaves+=1
        else:
            dfs(graph[s][i])

leaves = 0
if lst[target] == -1:
    print(0)
else:
    dfs(n)
    print(leaves)

생각 & 정리

리프를 세어주는 문제를 처음 만나게되어 당황했는데, 더 이상의 연결정보가 없으면 그 노드가 리프노드라는 것이 생각나 해결하게 되었다.