→ Problems

[Algorithm] 백준 - 1967번 트리의 지름 (Swift)

Swift librarian 2024. 5. 27. 11:53

문제 소개

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

트리에서 가장 긴 거리의 두 노드를 구하는 문제이다.
 

문제 풀이

아이디어는 다음과 같다. 한 노드에서 가장 먼 노드를 구하고, 그 노드에서 가장 먼 노드와의 거리를 구하면 된다.

import Foundation

let n = Int(readLine()!)!
var graph = [Int: [(node: Int, distance: Int)]]()

for i in 1...n {
    graph[i] = []
}

for _ in 0..<n-1 {
    let input = readLine()!.split(separator: " ").map { Int($0)! }
    let (root, node, distance) = (input[0], input[1], input[2])
    
    graph[root]?.append((node, distance))
    graph[node]?.append((root, distance))
}

var nod = 1
var dis = 0
var visited = Array(repeating: false, count: n+1)

func dfs(_ node: Int, _ cost: Int) {
    var next = false
    
    for i in graph[node]! {
        if !visited[i.node] {
            visited[i.node] = true
            next = true
            dfs(i.node, cost + i.distance)
        }
    }
    
    if !next && dis <= cost {
        dis = cost
        nod = node
    }
}

visited[nod] = true
dfs(nod, dis)

visited = Array(repeating: false, count: n+1)
visited[nod] = true
dfs(nod, 0)

print(dis)