[solved.ac 골드5] 1068_트리 (파이썬, DFS)

 

제한 사항

문제

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

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

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

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

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

입력

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

출력

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

예제

 

---

나의 풀이

골드 문제 치고는 어렵지 않았다 (?)

위 예제의 그림처럼 1번 노드를 지우면 1번 노드를 포함한 자식 노드들에 접근할 수 없다. 나는 여기서 바로 visited 리스트가 떠올랐다.

visited[1]을 True로 해버리면 1번 노드 이후부터는 접근하지 않으리라.

그리고 root 부터 DFS로 탐색하며 자식 노드의 개수가 0이면 answer 값을 하나씩 증가시켜 세어준다.

 

사용자 입력

import sys
input = sys.stdin.readline


n = int(input())

tree = [[] for _ in range(n)]
visited = [False] * n
root = -1
answer = 0

for node, parent in enumerate(map(int, input().split())):
    if parent == -1:
        root = node
    else:
        tree[parent].append(node)

visited[int(input())] = True

tree 리스트의 i번 인덱스에 저장되는 값은 i번 노드의 자식 노드 번호이다.

visited와 root, answer 변수를 선언한다.

이후 사용자에게 입력 받은 숫자들을 parent로 받는데 이때 enumerate를 사용하면 인덱스와 함께 받을 수 있다. 이 인덱스는 현재 노드를 뜻하게 된다.

parent가 -1이라면 루트 노드 이기에 root = node로 넣어준다.

그 외에는 tree[parent].append(node)로 자식 노드를 저장한다.

그리고 위에서 말했듯이 사용자에게 입력 받아 visited에서 True로 방문처리해준다.

 

DFS 함수 (틀린 코드)

def dfs(s):
    global answer

    if visited[s]: return
    visited[s] = True

    if len(tree[s]) == 0:
        answer += 1
        return

    for x in tree[s]:
        dfs(x)

s를 받아서 만약 방문했던 노드라면 빠져나오고, 그렇지 않다면 방문처리 하고 다음으로 넘어간다.

개수를 확인해 0이라면 자식이 없는 리프 노드라는 뜻이기에 answer를 증가시키고 빠져나온다.

만약 자식 노드가 있다면 새로운 DFS() 함수를 재귀적으로 호출한다.

 

그러나 위에 적어놨듯이 위 코드로는 틀렸다는 결과를 받게 되었다.

곰곰히 생각해보다가 알게된 사실은 다음과 같은 경우에 위 코드가 잘못 동작한다는 점이다.

내가 노드를 제거한 방식이 visited 리스트에서 방문 처리를 하는 방식으로 했기 때문에 tree[1]에는 여전히 [3] 이라는 값이 저장되어 있다.

위 상황에서 3번 노드가 제거되었을 때 1번 노드가 리프 노드가 되어야 하는데 이를 고려하지 못했던 것이다.

 

이를 수정한 코드는 다음과 같이 처리할 수 있겠다.

수정된 DFS 함수

# ...

deleted = int(input())
visited[deleted] = True

def dfs(s):
    global answer

    if visited[s]: return
    visited[s] = True

    if len(tree[s]) == 0:
        answer += 1
        return
    elif len(tree[s]) == 1 and tree[s][0] == deleted:
        answer += 1

    for x in tree[s]:
        dfs(x)

조건을 하나 추가했다.

자식이 하나이면서 그 자식이 deleted 된 노드라면 answer을 증가시킨다. 제거된 노드의 부모 노드가 리프 노드가 된 경우를 고려한 것이다.

 

전체 코드

import sys
input = sys.stdin.readline

n = int(input())

tree = [[] for _ in range(n)]
visited = [False] * n
root = answer = 0

for node, parent in enumerate(map(int, input().split())):
    if parent == -1:
        root = node
    else:
        tree[parent].append(node)

deleted = int(input())
visited[deleted] = True

def dfs(s):
    global answer

    if visited[s]: return
    visited[s] = True

    if len(tree[s]) == 0:
        answer += 1
        return
    elif len(tree[s]) == 1 and tree[s][0] == deleted:
        answer += 1

    for x in tree[s]:
        dfs(x)

dfs(root)
print(answer)

 

채점 결과