[BOJ] 1068 트리 (Python)

정말정말 오랜만에 쓰는 블로그 글!

원래는 알고리즘을 Swift로만 하다가 최근에는 파이썬도 하고 있다 

 

 

문제 링크 

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

 

문제 접근

트리에서 특정 노드를 지웠을 때, 남아있는 트리의 노드 중 리프 노드의 개수를 구하는 문제이다. 

 

트리 문제를 풀 때는 보통 그래프로 접근을 하고, 2차원 배열에 한 노드에 연결된 노드 번호를 모두 저장하는 식으로 구현을 한다. 

 

위의 그림으로 예를 들자면 아래처럼 저장을 한다는 이야기다. 

[
	[1,2],	  // 0번 노드
 	[0,3,4],  // 1번 노드
 	[0],	  // 2번 노드 
 	[1], 	  // 3번 노드
 	[1]	  	  // 4번 노드
]

 

하지만 리프 노드인지에 대한 판단을 어떻게 할 것인가? 라는 고민이 들었다.

그래서 생각한 방법은 아래와 같이 자식 노드만 저장하는 것이었다. 따라서 리스트에 값이 없으면 리프 노드라고 판단할 수 있다. 

[
	[1,2],	// 0번 노드
    [3,4],	// 1번 노드
    [],		// 2번 노드
    [],		// 3번 노드
    []		// 4번 노드
]

 

 

문제 풀이 

기본 입력 및 세팅

n = int(input())
parent = list(map(int, input().split()))
remove = int(input())

graph = [[] for _ in range(n)]	// 연결 정보 저장
root = 0	// 루트노드

 

 

루트 노드는 따로 `root` 변수에 저장해주고, 지울 노드인 `remove`를 제외하고 자식 노드만 `graph` 배열에 저장한다. 

사실 이렇게 하면 `remove`의 바로 아래 자식 노드만 제외되기 때문에 그건 아래에서 따로 처리를 해주려고 한다. 

for i in range(n):
	# 루트
    if parent[i] == -1: 
        root = i
        continue 
    # 지울 노드는 저장 x
    if parent[i] == remove or i == remove: continue 
    graph[parent[i]].append(i)

 

 

DFS 이용

DFS를 이용해서 리프 노드에 도달하면 리프 노드 개수를 1 증가시키고 반환하는 형태로 구현을 했다. `graph` 배열에는 자식 노드만 넣어놨으므로 현재 노드가 `cur`이라면 그 자식 노드는 `graph[cur]`에 담겨있다. 

answer = 0
visited = [False] * (n)
def dfs(cur):
    visited[cur] = True
    
    # 리프 노드
    if len(graph[cur]) == 0:
        global answer
        answer += 1
        return 

    for next in graph[cur]:
        if visited[next]: continue 
        dfs(next)

 

마지막으로 dfs 함수는 루트부터 시작하도록 호출을 하는데, 루트 노드가 지워야 할 `remove`와 같다면 dfs 함수를 호출하지 않아서 답이 0 으로 나오도록 한다.

if root != remove:
    dfs(root)

 

 

전체 코드

import sys

input = sys.stdin.readline
n = int(input())
parent = list(map(int, input().split()))
remove = int(input())
graph = [[] for _ in range(n)]
root = 0

for i in range(n):
	# 루트 노드
    if parent[i] == -1: 
        root = i
        continue 
    # 지울 노드
    if parent[i] == remove or i == remove: continue 
    graph[parent[i]].append(i)

# DFS
answer = 0
visited = [False] * (n)
def dfs(cur):
    visited[cur] = True
    # 리프 노드 
    if len(graph[cur]) == 0:
        global answer
        answer += 1
        return 

    for next in graph[cur]:
        if visited[next]: continue 
        dfs(next)

if root != remove:
    dfs(root)

print(answer)