[백준] 1068 - 트리 (Python)

문제 링크

알고리즘

• 그래프 탐색

Tip

• 입력으로 주어진 노드 중 부모가 -1인 노드가 루트 노드

풀이

1. 입력을 받아서 부모가 -1인 노드를 root_node 변수에 저장한다.
2. 삭제할 노드가 아닌 노드로 그래프를 만든다.

1, 2번을 동시에 진행한다.

import sys

input = sys.stdin.readline

N = int(input())


# 부모 -> 자식
graph = [[] for _ in range(N)]  # graph[i] -> i 번째 노드의 자식 노드들 번호 리스트
parents = list(map(int, input().split())) # 부모 노드
root_node = 0 # 루트 노드
delete_node = int(input()) # 삭제할 노드


for i in range(N):
    if parents[i] == -1: # 부모 노드가 -1이면 루트 노드
        root_node = i
    else:
        # 삭제할 노드가 아닌 경우만 부모의 자식 리스트에 추가
        if i != delete_node:
            graph[parents[i]].append(i)

2. 삭제할 노드가 루트 노드라면 0을 출력하고 프로그램을 종료한다.

if delete_node == root_node:
    print(0)
    exit(0)

3. 모든 노드를 dfs로 순회하면서 자식 노드가 없는 노드를 발견하면 결과에 1을 추가한다.

순회가 끝나면 결과를 출력한다.

res = 0 # 리프노드의 개수


def dfs(node: int):
    global res

    # 자식이 없으면 리프노드
    if len(graph[node]) == 0:
        res += 1
        return
    
    # 자식 노드들을 dfs로 재귀 호출
    for n in graph[node]:
        dfs(n)


dfs(root_node)

print(res)

전체 코드

import sys

input = sys.stdin.readline

N = int(input())


# 부모 -> 자식
graph = [[] for _ in range(N)]  # graph[i] -> i 번째 노드의 자식 노드들 번호 리스트
parents = list(map(int, input().split())) # 부모 노드
root_node = 0 # 루트 노드
delete_node = int(input()) # 삭제할 노드


for i in range(N):
    if parents[i] == -1: # 부모 노드가 -1이면 루트 노드
        root_node = i
    else:
        # 삭제할 노드가 아닌 경우만 부모의 자식 리스트에 추가
        if i != delete_node:
            graph[parents[i]].append(i)

if delete_node == root_node:
    print(0)
    exit(0)


res = 0


def dfs(node: int):
    global res

    # 자식이 없으면 리프노드
    if len(graph[node]) == 0:
        res += 1
        return
    for n in graph[node]:
        dfs(n)


dfs(root_node)

print(res)