[Python] 백준 1068 트리

문제

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

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

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

 

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

 

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

입력

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

출력

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

예제 입력 1 복사

5
-1 0 0 1 1
2

예제 출력 1 복사

2

예제 입력 2 복사

5
-1 0 0 1 1
1

예제 출력 2 복사

1

예제 입력 3 복사

5
-1 0 0 1 1
0

예제 출력 3 복사

0

예제 입력 4 복사

9
-1 0 0 2 2 4 4 6 6
4

예제 출력 4 복사

2

 

---

코드

import sys
sys.setrecursionlimit(10**6)

def removeTreeDfs(A, removeIndex):
    A[removeIndex] = None
    for i in range(len(A)):
        if removeIndex == A[i]:
            removeTreeDfs(A, i)



N = int(sys.stdin.readline().rstrip()) # 노드개수
nodeList = list(map(int,sys.stdin.readline().rstrip().split())) # 노드 리스트
trueList = [True for i in range(N)]
removeIndex = int(sys.stdin.readline().rstrip()) # 삭제 인덱스

# remove subTree
removeTreeDfs(nodeList, removeIndex)

# find leaf node
leafCnt = 0
for i in range(N):
    if nodeList[i] != None and i not in nodeList:
        leafCnt += 1

print(leafCnt)

---

리뷰

하나의 노드를 삭제한 후의 리프 노드의 수를 구하는 문제였다.

트리의 노드 삭제 연산을 DFS를 이용해 진행하였다.

 

첫 시도에서 실수한 점은 트리가 무조건 이진트리라고 가정하고 했다는 것이었다.

따라서 첫 childIdx를 찾고 그것의 +1까지만 고려하였는데

예시에서는 문제가 없었지만 실제 제출한 코드는 틀리게 되었다.

 

따라서 모든 자식 노드를 찾을 수 있도록 다시 바꾸어서 코드를 작성하고

정답을 받았다.

def removeTreeDfs(A, removeIndex):
    A[removeIndex] = None
    # 자식 노드 찾기
    for i in range(len(A)):
        if removeIndex == A[i]:
            removeTreeDfs(A, i)