백준 - LCA 2 (11438)

import Foundation

// FileIO ..

let file = FileIO()
let N = file.readInt()
var tree: [[Int]] = Array(repeating: [], count: N+1)
// 최대 k값 구하기
let kmax = Int(floor(log2(Double(N))))

var parent = Array(repeating: Array(repeating: 0, count: N + 1), count: kmax + 1)

// 트리 만들기
for _ in 0..<N - 1 {
    let a = file.readInt()
    let b = file.readInt()
    tree[a].append(b)
    tree[b].append(a)
}

// 깊이 구하기
var depth = Array(repeating: 0, count: N + 1)
var visited = Array(repeating: false, count: N + 1)

dfs(1, 0)

func dfs(_ node: Int, _ dep: Int) {
    visited[node] = true
    depth[node] = dep
    for next in tree[node] {
        if !visited[next] {
            parent[0][next] = node
            dfs(next, dep + 1)
        }
    }
}

// 부모 정보 저장
for k in 1...kmax {
    for n in 1...N {
        parent[k][n] = parent[k - 1][parent[k - 1][n]]
    }
}

// 질의 수행
let M = file.readInt()

for _ in 0..<M {
    let a = file.readInt()
    let b = file.readInt()
    let result = findLCA(a, b)
    print(result)
}

func findLCA(_ num1: Int, _ num2: Int) -> Int {
    var a = num1
    var b = num2
    if depth[a] > depth[b] {
        (a, b) = (b, a)
    }
    
    for k in (0...kmax).reversed() { // a와 b의 깊이 맞추기
        if Int(pow(2.0, Double(k))) <= depth[b] - depth[a] { // 깊이 차이가 충분하다면
            if depth[a] <= depth[parent[k][b]] {
                b = parent[k][b]
            }
        }
    }
    // 조상 빠르게 찾기
    for k in (0...kmax).reversed() {
        if parent[k][a] != 0 && parent[k][a] != parent[k][b] {
            a = parent[k][a]
            b = parent[k][b]
        }
    }
    var LCA = a
    if a != b {
        LCA = parent[0][LCA]
    }
    return LCA
}

질의의 수(M)rk 100,000으로 크기 때문에 일반적인 LCA 알고리즘을 개선한 방법을 사용해야 한다.

입력값이 많은 탓에 채점이 느리게 되어 파일입출력 코드를 사용해서 문제를 풀었다.

실수한 점

분명 맞게 풀었는데 자꾸 처음부터 WA가 떠서 코드를 붙잡고 계속 고민해봤다.

알고보니 LCA 함수를 실행할 때 원래는 두 수의 깊이를 비교해서 깊이가 작은 값을 a로 두어야 하는데 나는 depth 값이 아니라 a, b 자체의 값을 비교해버려서 틀렸다고 나온 것이었다.

저번에도 이렇게 실수했던 것 같은데.. 제발 헷갈리지 말자!!

LCA 알고리즘 역시 한 번에 이해하기 어려운 개념이었다. 개념도 꼭 나중에 다시 복습하고 구현하는 법에도 익숙해지도록 해야겠다.