[백준] 1068번: 트리(Python)

문제 링크

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

1068번: 트리

문제 풀이

풀이

• 재귀함수를 이용해서, 지우려는 노드를 실제로 dict 형태를 가진 Tree에서 삭제시키는 것.
• 삭제 진행
  • 자식 길이가 0이면 해당 노드 삭제
  • 돌아와서 본인 삭제
• 삭제 과정 후에 자식이 0인 것을 센다.

import sys
input = sys.stdin.readline
sys.setrecursionlimit(100000)
N = int(input())
Tree = {}
for i in range(N):
    Tree[i] = []

first = 0 # top root
parent = list(map(int, input().split()))
re_node = int(input())

for i, v in enumerate(parent):
    if v != -1:#맨 위 root 노드가 아닌 것으로
        Tree[v].append(i)

def remove(root):
    for node in Tree[root]:
        if not Tree[node]: #자식 노드가 없는 것부터 제거한다.
            Tree.pop(node)
        else:# 자식 노드가 있으면 재귀함수로 나올 때까지
            remove(node)
    Tree.pop(root) # 자식들 다 없애고 마지막에 root 노드까지 제거

remove(re_node) # 지울 지점부터


cnt = 0
for key in Tree.keys():
    if re_node in Tree[key] and len(Tree[key])-1 == 0: #지우려 한 노드를 자식으로 둔 것은 빼고 생각하기.
        # Tree[key].pop(Tree[key].index(re_node))
        cnt += 1
    if not Tree[key]:
        cnt += 1

print(cnt)

참고한 코드 풀이

[참고 링크]

• 각 노드의 부모 번호를 알려준 입력 값이므로 이를 활용한 풀이이다. 
• 삭제할 노드 배열 값을 삭제하는 의미로 -2로 바꾼다. (-1은 루트노드와 겹치므로 피하기)
  • 이후 배열 전체를 탐색 -> 방금 삭제한 인덱스를 부모 노드로 가진 노드로 다시 재귀함수 돌리기
• 재귀가 끝나면 삭제될 노드는 전부 -2 로 갱신되어 있다. 
• -2가 아니고, 다른 노드의 부모 노드도 아닌 원소를 찾을 때 count += 1 을 해준다. 

import sys
input = sys.stdin.readline

def dfs(num, arr):
    arr[num] = -2
    for i in range(len(arr)):
        if num == arr[i]:
            dfs(i, arr)

n = int(input())
arr = list(map(int, input().split()))
k = int(input())
count = 0

dfs(k, arr)
for i in range(len(arr)):
    if arr[i] != -2 and i not in arr:
        count += 1
print(count)

 

후기

다 풀고 다른 코드를 보다가 짧은 코드여서 넣어봤는데, 성능은 매번 재귀함수마다 전체 배열을 탐색해야해서 조금 느려진듯하다.

아래가 내 코드, 위에가 짧은 참고 코드를 넣었을 때 실행한 결과이다.