[BOJ - DFS] 1068번 : 트리 ( python )

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

1068번: 트리

 

 

문제

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

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

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

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

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

입력

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

출력

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

 

 

---

 

1. 사고하기

• 목표 : 리프 노드의 개수
• 유형 : 트리하면 그래프 > 깊이우선탐색 이라고 생각한다.
• dfs함수로 들어온 노드 v는 삭제되는 노드이므로 tree[v] 를 -2로 바꾸어준다. (tree의 범위는 -1~ 50이므로)
• 노드 개수만큼 반복해서 index로 접근, tree[i] 가 v이면 삭제해야하는 노드의 자식이 i라는 뜻 > i의 자식도 지워주어야하므로 dfs(i)해준다. 

>> dfs 끝난 이후엔 tree에서 삭제되어야 하는 노드번호들이  -2로 변경되어 있다. 

n 만큼 반복해서 tree를  체크해보자.

 

리프 노드의 조건은 2가지가 있다. 이 2가지에 해당하면 cnt를 늘려주고 마지막에 cnt를 출력해주면 된다.

>>>> 1번 : tree[i] 가 -2이면 안된다 (지운노드이거나 지운노드의 자식이거나..)

>>>> 2번 : tree에는 부모의 정보가 들어있으므로 i라는 값이 tree안에 있으면 자식이 있는 노드이다. 따라서 

i not in tree 를 사용해 체크한다. 

 

2. 이 문제를 공유한 이유 >> 테스트케이스는 다맞았는데 틀렸다고 나와서 당황한 문제

초안 코드를 작성했을 땐 테스트케이스가 다 맞았는데 바로 틀렸습니다가 나와서 예외처리를 어떻게 해줘야할지 막막했다.  

검색을 해보니 어떤 노드를 지웠을 때 그 노드가 부모와 자식이 각각 한개밖에 없는 경우엔 노드를 지워도 리프는 1개가 남았다나... 

 

조건 넣는게 너무 어려워서 아예 코드를 다 갈아 엎었다. 

나중에 조건 넣는걸로 다시 풀어봐야겠다고 생각이 들었다. 

 

정답코드가 깔끔하긴 한데 초안코드 포기를 못하겠다 ㅠ 

 

초안 코드

import sys
input = sys.stdin.readline

def dfs(v):
    visited[v] = -2
    for i in graph[v]:
        dfs(i)
            
n = int(input())
li = list(map(int, input().split()))

graph = [[] for _ in range(n)]
#그래프는 자식을 담은 리스트로서 사용
# graph [[1,2],[3,4],[],[],[]]
for i in range(1,n):
    graph[li[i]].append(i)

erase = int(input())
visited = [0] * n

dfs(erase)

cnt = 0
for i in range(n):
    if visited[i] != -2 and len(graph[i]) != 0:
        cnt+=1

print(cnt)

 

 

 

 

정답 코드 

#1068_트리_dfs_gold5

import sys
input = sys.stdin.readline

def dfs(v):
    tree[v] = -2
    for i in range(n): #전체 트리 반복
        if v == tree[i]: # 지우고자 하는 노드 v가  tree[i]에 들어있으면 tree[i]는 v의 자식
            dfs(i) # i의 자식도 지움
            
n = int(input())
tree = list(map(int, input().split())) # [-1, 0, 0, 1, 1]
erase = int(input())

dfs(erase)
cnt = 0

for i in range(n):
    if tree[i] != -2 and i not in tree: #-2는 지우는 노드들 // i노드의 자식이 트리 안에 없으면 == 리프노드임
        cnt+=1

print(cnt)