트리에서의 거리

1. 문제

2. 입출력

3. 입출력 예시

4. 문제 설계

• 특정 노드에서부터 시작해서 깊이 우선 탐색을 진행하게 되면, 특정 노드에서 다른 모든 노드까지의 거리를 알 수 있음
• 따라서 트리를 구성할 때 모든 부모 자식 관계의 노드는 서로 양방향으로 이어주고, X 노드에서 부터 깊이 우선 탐색을 진행하면서
• 별도의 배열에 거리를 저장함으로써 문제를 해결할 수 있음

5. 전체 코드

//MARK: - 트리에서의 거리

//MARK: - Framework
import Foundation

//MARK: - Type
struct Node {
    //MARK: - Property
    private var edges: [Int]
    
    //MARK: - initializer
    init() {
        self.edges = []
    }
    
    //MARK: - Method
    mutating func push(_ data: Int) -> Void {
        self.edges.append(data)
    }
    
    func getNodeData(_ index: Int) -> Int {
        return self.edges[index]
    }
    
    func size() -> Int {
        return self.edges.count
    }
}

//MARK: - Function
func dfs(_ tree: [Node], _ visited: inout [Bool], _ levels: inout [Int], _ node: Int, _ currentLevel: Int) -> Void {
    visited[node] = true
    levels[node] = currentLevel
    
    for i in 0..<tree[node].size() {
        let y: Int = tree[node].getNodeData(i)
        
        if !visited[y] {
            dfs(tree, &visited, &levels, y, currentLevel + 1)
        }
    }
}

func solution() -> Void {
    //MARK: - Input
    guard let input: [String] = readLine()?.components(separatedBy: " ") else { return }
    
    let n: Int = input.map { Int($0) }[0] ?? 0
    let X: Int = input.map { Int($0) }[1] ?? 0
    let Y: Int = input.map { Int($0) }[2] ?? 0
    
    var tree: [Node] = Array(repeating: Node(), count: n + 10)
    var visited: [Bool] = Array(repeating: false, count: n + 10)
    var levels: [Int] = Array(repeating: 0, count: n + 10)
    var distance: Int = 0
    
    for _ in 0..<n - 1 {
        guard let inputData: [String] = readLine()?.components(separatedBy: " ") else { return }
        
        let a: Int = inputData.map { Int($0) }[0] ?? 0
        let b: Int = inputData.map { Int($0) }[1] ?? 0
        
        tree[a].push(b)
        tree[b].push(a)
    }
    
    //MARK: - Process
    dfs(tree, &visited, &levels, X, 0)
    distance = levels[Y]
    
    //MARK: - Output
    print(distance)
}
solution()

 

전체코드는 여기에서 확인할 수 있습니다.