1068 - 트리

info

• 문제 보기: 1068 - 트리
• 소요 시간: 49분
• 풀이 언어: python
• 체감 난이도: 2️⃣~3️⃣
• 리뷰 횟수: ✅

풀이 키워드

스포주의

자료구조 트리 dfs

풀이 코드

info

• 메모리: 31120 KB
• 시간: 40 ms

import sys
input = sys.stdin.readline

children = []

def dfs(cur, target):
    cnt = 0
    if cur == target:
        return 0
    if not children[cur]: # leaf
        return 1
    for nxt in children[cur]:
        cnt += dfs(nxt, target)
    
    if not cnt: # not leaf but none found -> just became leaf
        cnt += 1
    return cnt

def solution():
    global children
    n = int(input())
    data = [int(x) for x in input().rstrip().split(' ')]
    children = [[] for _ in range(n)]
    root = -1
    for i, parent in enumerate(data):
        if ~parent:
            children[parent].append(i)
        else:
            root = i
    target = int(input())
    print(dfs(root, target))

    
if __name__ == '__main__':
    solution()

풀이 해설

입력값들을 적절히 구조화하고, 이를 dfs로 스캔하여 답을 계산하는 방식이다.

caution

응 아니야
루트 노드의 인덱스가 반드시 0이라는 보장이 없다.

for i, parent in enumerate(data):
    if ~parent:
        children[parent].append(i)
    else:
        root = i

그러므로 ~parent != 0인 경우에 대해서만 자식 노드 인덱스를 기록하고,

~parent가 0이 되는 경우에 대해서 루트 노드 처리한다.

📌DFS

def dfs(cur, target):
    cnt = 0
    if cur == target:
        return 0
    if not children[cur]: # leaf
        return 1
    for nxt in children[cur]:
        cnt += dfs(nxt, target)
    
    if not cnt: # not leaf but none found -> just became leaf
        cnt += 1
    return cnt

입력값을 구조화했으면 dfs를 돌리는데, return이 일어나는 edge case를 신경 쓸 필요가 있다.

1. if cur == target: 삭제할 노드인 경우 그 밑의 depth로 더이상 재귀를 진행하지 않는다. --> +0
2. if not children[cur]: 자식 노드가 없는 노드이므로, leaf이다. --> +1
3. if not cnt: 이전 라인에서 자식 노드를 모두 갔다왔음에도 자식 노드 수가 0으로 찍히는 경우. 이게 무슨 소리지? 바로 모든 자식이 1번에서 return 0 당한 case이다. (사실상 노드 삭제는 1개만 이루어지므로, 자식이 1개뿐이었는데 얘가 삭제된 것) 이 경우 자기 자신이 leaf가 된다. --> +1

메모

• 루트 노드 낚시는 다음에도 또 낚일 것 같음
• 생각보다 list 연산 오버헤드가 별로 들지 않아서 의외였던 문제