[백준] 1967번 트리의 지름 [Swift]

문제 링크

1967번: 트리의 지름

문제

트리(tree)는 사이클이 없는 무방향 그래프이다. 트리에서는 어떤 두 노드를 선택해도 둘 사이에 경로가 항상 하나만 존재하게 된다. 트리에서 어떤 두 노드를 선택해서 양쪽으로 쫙 당길 때, 가장 길게 늘어나는 경우가 있을 것이다. 이럴 때 트리의 모든 노드들은 이 두 노드를 지름의 끝 점으로 하는 원 안에 들어가게 된다.

이런 두 노드 사이의 경로의 길이를 트리의 지름이라고 한다. 정확히 정의하자면 트리에 존재하는 모든 경로들 중에서 가장 긴 것의 길이를 말한다.

입력으로 루트가 있는 트리를 가중치가 있는 간선들로 줄 때, 트리의 지름을 구해서 출력하는 프로그램을 작성하시오. 아래와 같은 트리가 주어진다면 트리의 지름은 45가 된다.

입력

파일의 첫 번째 줄은 노드의 개수 n(1 ≤ n ≤ 10,000)이다. 둘째 줄부터 n-1개의 줄에 각 간선에 대한 정보가 들어온다. 간선에 대한 정보는 세 개의 정수로 이루어져 있다. 첫 번째 정수는 간선이 연결하는 두 노드 중 부모 노드의 번호를 나타내고, 두 번째 정수는 자식 노드를, 세 번째 정수는 간선의 가중치를 나타낸다. 간선에 대한 정보는 부모 노드의 번호가 작은 것이 먼저 입력되고, 부모 노드의 번호가 같으면 자식 노드의 번호가 작은 것이 먼저 입력된다. 루트 노드의 번호는 항상 1이라고 가정하며, 간선의 가중치는 100보다 크지 않은 양의 정수이다.

문제 풀이

트리의 지름이라는 개념이 생소해서 신선한 문제였던거 같습니다.

우선 이 문제에서 주어진 그림처럼 트리의 지름을 구하기 위해서는 2개의 정점이 필요한 것을 알 수 있었습니다!

하나의 점은 트리 속 정점중 어떤 점을 선택했을 때 모든 정점에 도달하는 가중치가 가장 큰 정점입니다.

 

이러한 것은 지름에 해당하는 정점을 선택하더라도 성립됩니다!

 

이렇게 지름에 해당하는 두 개의 점 중 하나를 찾았다면 그 점에서 모든 정점에 도달하는 가중치를 구한 뒤 가장 큰 값이 정답이 됩니다.

 

따라서 이 문제를 풀기 위해선 임의의 정점에서 지름을 구성하는 두 점 중 하나를 찾기위한 탐색 1번,

찾은 하나의 점에서 다른 하나의 점을 찾기위한 탐색 1번을 해주시면 됩니다.

즉 트리를 모두 탐색하는 dfs, bfs 중 하나를 2번 사용해주시면 됩니다.

저는 dfs를 사용하여 문제를 풀었습니다!

 

처음 탐색에서는 트리의 어떤점에서 탐색을 시작하더라도 상관이 없습니다.

두 번째 탐색에서는 첫 탐색에서 구한 하나의 정점에서 탐색을 시작해주시면 됩니다!

 

dfs, bfs를 2번 사용해야한다는 점과 단순히 방문만 하는 것이 아닌 가중치를 계속 기억해줘야한다는 점이 재밌었던 문제입니다!

 

소스 코드

func solution() {
    let n = Int(readLine()!)!
    
    // 간선의 정보를 저장한 Array
    var graph: [[(Int,Int,Int)]] = [[(Int,Int,Int)]](repeating: [], count: n+1)
    
    // 입력 받는 중
    for _ in 0..<n-1 {
        let line = readLine()!.split(separator: " ").map({Int(String($0))!})
        graph[line[0]].append((line[0],line[1],line[2]))
        graph[line[1]].append((line[1],line[0],line[2]))
    }
    
    // 어떤 정점에 가기 위한 가중치를 저장하는 Array. 초기 값을 0으로 설정.
    var visited: [Int] = [Int](repeating: 0, count: n+1)
    
    // 트리의 정점을 모두 확인하기 위한 dfs
    func dfs (_ start: (Int,Int,Int)) {
        var stack: [(Int,Int,Int)] = [start]
        
        while !stack.isEmpty {
            let now = stack.popLast()!
            
            if visited[now.1] == 0 {
                visited[now.1] = visited[now.0] + now.2
                
                stack += graph[now.1]
            }
        }
    }
    
    // 트리의 어떤 정점에서 시작했을 때 가중치가 가장 큰 값이 트리의 지름을 구성하는 두 점 중 하나.
    dfs((3,3,0))
    let max = visited.max()!
    var temp = 0
    
    for i in 0..<visited.count {
        if visited[i] == max {
            temp = i
            break
        }
    }
    // 구해진 한점을 기준으로 모든 정점으로 가는 가중치를 구한다.
    visited = [Int](repeating: 0, count: n+1)
    dfs((temp,temp,0))
    // 정점에 도달하는 가중치의 최댓값이 정답
    print(visited.max()!)
}

solution()