1068. 트리 (Swift, Python)

1068번: 트리

 

풀이

 2차원 배열을 만들어 각 노드에 자신의 자식 노드 리스트를 담아둔다. 예를 들어 입력이 -1 0 0 1 1 과 같은 형태로 주어지면 0은 root 노드로 설정하고 1과 2는 부모 노드가 0이므로 배열의 [0]에 [1, 2]를 담는다. 3과 4는 부모 노드가 1이므로 배열의 [1]에 [3, 4]를 담는다.

 

 그 후 root 노드부터 dfs 탐색을 시작한다. dfs로 노드를 탐색하면서 현재 노드에 자식 노드가 하나도 없을 경우 카운트를 증가시키고 탐색을 종료한다. 자식 노드 중에 지울 노드가 있을 경우에는 dfs를 돌지 않는데 하나 주의할 것은 현재 노드에 지워야 하는 노드 하나만 있는 경우다. 이럴 때는 따로 카운트를 증가시켜야 한다.

 이러한 경우는 아래의 입력 경우를 생각해보자.

5
-1 0 1 2 3 
2

 

코드

[Swift]

// 트리 dfs 탐색
func dfs(_ node: Int) {
    
    // 현재 node의 자식이 하나도 없을 경우 리프 노드이므로 res 1 증가
    if nodes[node].count == 0 {
        res += 1
        return
    }
    
    // node의 자식 노드 for문
    for i in nodes[node] {
        
        // 자식 노드가 지워야 하는 노드일 경우 dfs를 호출하지 않음
        // 지워야 하는 노드가 유일한 자식 노드일 경우에는 현재 node가 리프 노드가 되므로 res 1 증가
        if i == M {
            if nodes[node].count == 1 {
                res += 1
            }
            continue
        }
        
        dfs(i) // 재귀 호출
    }
}
 
var N: Int = 0 // 트리의 노드의 개수
var M: Int = 0 // 지울 노드 번호
var p: [Int] = [] // 각 노드의 부모를 담아줄 배열
var root: Int = 0 // 루트 노드(부모 노드가 -1인 노드)
var res: Int = 0 // 리프 노드의 총 개수
 
if let input = readLine() {
    N = Int(input)!
}
 
var nodes:[[Int]] = Array(repeating: [Int](), count: N) // [node]의 자식 노드 리스트
 
if let input = readLine() {
    p = input.split(separator: " ").map { Int($0)! }
}
 
if let input = readLine() {
    M = Int(input)!
}
 
for i in 0..<N {
    if p[i] == -1 {
        root = i // 루트 노드 설정
        continue
    }
    
    nodes[p[i]].append(i)
}
 
// 루트 노드가 지워야 하는 노드일 경우에는 dfs 탐색을 하지 않음
if M != root {
    dfs(root)
}
 
print(res)
 

[Python]

# 트리 dfs 탐색
def dfs(node):
    global res
 
    # 현재 node의 자식이 하나도 없을 경우 리프 노드이므로 res 1 증가
    if len(p[node]) == 0:
        res += 1
        return
 
    # node의 자식 노드 for문
    for i in p[node]:
 
        # 자식 노드가 지워야 하는 노드일 경우 dfs를 호출하지 않음
        # 지워야 하는 노드가 유일한 자식 노드일 경우에는 현재 node가 리프 노드가 되므로 res 1 증가
        if i == M:
            if len(p[node]) == 1:
                res += 1
            continue
 
        dfs(i) # 재귀 호출
 
 
if __name__ == '__main__':
    N = int(input()) # 트리의 노드의 개수
    nodes = list(map(int, input().split())) # 각 노드의 부모를 담아줄 배열
    M = int(input()) # 지울 노드 번호
    res = 0 # 리프 노드의 총 개수
 
    p = [[] for _ in range(N)] # [node]의 자식 노드 리스트
    root = 0 # 루트 노드(부모 노드가 -1인 노드)
 
    for i in range(N):
        if nodes[i] == -1:
            root = i # 루트 노드 설정
        else:
            p[nodes[i]].append(i)
 
    # 루트 노드가 지워야 하는 노드일 경우에는 dfs 탐색을 하지 않음
    if root != M:
        dfs(root)
 
    print(res)