백준 - 트리 (1068)

첫 번째 풀이

let N = Int(readLine()!)!
var tree: [[Int]] = Array(repeating: [], count: N)
var visited = Array(repeating: false, count: N)
var answer = 0

// 트리 만들기
let input = readLine()!.split(separator: " ").map { Int(String($0))! }
var root = 0
for i in 0..<N {
    let parent = input[i]
    if parent == -1 { // 루트 노드인 경우
        root = i
    } else { // 아니면 트리(인접 리스트)에 저장
        tree[parent].append(i)
    }
}

let target = Int(readLine()!)!

// 지울 노드가 루트 노드가 아니면 DFS 실행
if target != root {
    dfs(root)
}

print(answer)

func dfs(_ node: Int) {
    visited[node] = true
    if tree[node].isEmpty {
        answer += 1
        return
    }
    for next in tree[node] {
        if !visited[next] && next != target {
            dfs(next)
        }
    }
}

입력받은 노드들을 인접 리스트에 저장한 후, DFS 알고리즘을 실행하며 리프 노드를 찾는다. 이때 다음 노드가 제거할 노드면 DFS를 실행하지 않는다.

이 코드는 78% 즈음에서 WA가 떴다. 먼저 반례는 아래와 같다.

2
-1 0
1

정답 : 1

9
-1 0 0 5 2 4 4 6 6
4

정답 : 2

첫 번째 반례만 보고 잘못된 점을 금방 찾을 수 있었다. 위 코드에서는 탐색 중인 노드가 지금 리프 노드(자식 노드가 없는 경우)여야만 answer 값에 1을 추가시킨다.

그래서 첫 번째 반례처럼 자식 노드를 지워서 현재 노드가 리프 노드가 되는 경우는 포함을 하지 못한다.

최종 풀이

let N = Int(readLine()!)!
var tree: [[Int]] = Array(repeating: [], count: N)
var visited = Array(repeating: false, count: N)
var answer = 0

// 트리 만들기
let input = readLine()!.split(separator: " ").map { Int(String($0))! }
var root = 0
for i in 0..<N {
    let parent = input[i]
    if parent == -1 { // 루트 노드인 경우
        root = i
    } else { // 아니면 트리(인접 리스트)에 저장
        tree[parent].append(i)
    }
}

let target = Int(readLine()!)!

// 지울 노드가 루트 노드가 아니면 DFS 실행
if target != root {
    dfs(root)
}

print(answer)

func dfs(_ node: Int) {
    visited[node] = true
    var hasChild = false
    for next in tree[node] {
        if !visited[next] && next != target {
            dfs(next)
            hasChild = true
        }
    }
    if !hasChild {
        answer += 1
    }
}

일단 for문을 돌려서 자식 노드가 존재하면 hasChild 의 값을 true 로 설정한다. hasChild 의 값이 false 이면 더 이상 탐색할 자식 노드가 없는 리프 노드이므로 이때 answer 의 값을 증가시킨다.