[백준] 2887번 행성 터널 [Swift]

문제 링크

2887번: 행성 터널

문제

때는 2040년, 이민혁은 우주에 자신만의 왕국을 만들었다. 왕국은 N개의 행성으로 이루어져 있다. 민혁이는 이 행성을 효율적으로 지배하기 위해서 행성을 연결하는 터널을 만들려고 한다.

행성은 3차원 좌표 위의 한 점으로 생각하면 된다. 두 행성 A(xA, yA, zA)와 B(xB, yB, zB)를 터널로 연결할 때 드는 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이다.

민혁이는 터널을 총 N-1개 건설해서 모든 행성이 서로 연결되게 하려고 한다. 이때, 모든 행성을 터널로 연결하는데 필요한 최소 비용을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 행성의 개수 N이 주어진다. (1 ≤ N ≤ 100,000) 다음 N개 줄에는 각 행성의 x, y, z좌표가 주어진다. 좌표는 -109보다 크거나 같고, 109보다 작거나 같은 정수이다. 한 위치에 행성이 두 개 이상 있는 경우는 없다. 

출력

첫째 줄에 모든 행성을 터널로 연결하는데 필요한 최소 비용을 출력한다.

문제 풀이

이 문제는 무방향 가중치 그래프에서 모든 정점을 최소 비용으로 연결하는 최소 스패닝 트리 문제입니다!

가장 간단한 방법인 모든 행성에 대한 간선을 만들어 버리면 메모리 초과는 물론 시간 초과도 발생하게 되기 때문에 다른 접근방법이 필요해요...

 

우선 이 문제에서 사실상 x, y, z 좌표가 주어져 있긴 하지만 실제 두 점 사이의 거리를 구하는 방식으로 가중치를 구하는 것이 아닌 그냥 x, y, z값 중 가장 작은 하나의 값이 가중치가 되기 때문에 (문제에서 주어진 조건인 min(|xA-xB|, |yA-yB|, |zA-zB|) 때문입니다.) 아주 재밌는 접근이 가능했습니다.

 

바로 x,y,z 값을 기준으로 한 번씩 정렬한 뒤 인접한 행성과의 거리를 간선의 후보로 두는 것이죠.

예를 들면 입력 예제로 주어진 것으로 보겠습니다.

0번 행성	1번 행성	2번 행성	3번 행성	4번 행성
11, -15, -15	14, -5, -15	-1, -1, -5	10, -4, -1	19, -4, 19

이렇게 주어져 있습니다. 이렇게 주어진 행성들을 먼저 x값 기준으로 정렬하는 거예요

2번 행성	3번 행성	0번 행성	1번 행성	4번 행성
-1, -1, -5	10, -4, -1	11, -15, -15	14, -5, -15	19, -4, 19

이렇게 하면 간선의 후보가 될 수 있는 녀석들을 인접한 행성들, 즉 2,3번 행성(가중치 : 11), 3,0번 행성(가중치 : 1), 0,1번 행성(가중치 : 3), 1,4번 행성(가중치 : 5)에서 x값을 가중치로 갖는 간선들이 후보로 추가됩니다.. 그런 뒤 이번에는 y값 기준으로 정렬을 진행합니다.

0번 행성	1번 행성	3번 행성	4번 행성	2번 행성
11, -15, -15	14, -5, -15	10, -4, -1	19, -4, 19	-1, -1, -5

이번에는 0,1번 행성(가중치 : 10), 1,3번 행성(가중치 : 1), 3,4번 행성(가중치 : 0), 4,2번 행성(가중치 : 3)에서 y값을 가중치로 갖는 간선들이 후보로 추가됩니다.

마지막으로 z값 기준으로 정렬을 진행합니다.

0번 행성	1번 행성	2번 행성	3번 행성	4번 행성
11, -15, -15	14, -5, -15	-1, -1, -5	10, -4, -1	19, -4, 19

이번에는 0,1번 행성(가중치 : 0), 1,2번 행성(가중치 : 10), 2,3 행성(가중치 : 4), 3,4 행성(가중치 : 20)에서 z값을 가중치로 갖는 간선들이 후보로 추가됩니다.

 

이렇게 한 뒤 지금까지 후보로 만들어둔 간선들로 크루스칼 알고리즘을 수행해주면 답이 나온답니다!

재밌는 문제였던 거 같습니다.

소스 코드

import Foundation

func solution(){
    let n = Int(readLine()!)!
    var parent = Array(0..<n)
    var planets: [[Int]] = []
    var distanceData: [(Int,Int,Int)] = []
    var lines: Int = 0
    var result: Int = 0
    
    func find(_ index: Int) -> Int{
        if parent[index] == index{
            return index
        } else {
            parent[index] = find(parent[index])
            return parent[index]
        }
    }
    func union(_ a: Int, _ b: Int){
        let p1 = find(a)
        let p2 = find(b)
        if p1 < p2 {
            parent[p2] = p1
        } else {
            parent[p1] = p2
        }
    }
    
    for i in 0..<n{
        var data = readLine()!.split(separator: " ").map({Int(String($0))!})
        data.append(i)
        planets.append(data)
    }
    
    // x값 기준 정렬
    planets.sort { (a, b) -> Bool in
        return a[0] < b[0]
    }
    
    for i in 0..<planets.count-1{
        // 행성번호1, 행성번호2, 가중치
        distanceData.append((planets[i][3],planets[i+1][3],abs(planets[i][0]-planets[i+1][0])))
    }
    
    // y값 기준 정렬
    planets.sort { (a, b) -> Bool in
        return a[1] < b[1]
    }
    
    for i in 0..<planets.count-1{
        // 행성번호1, 행성번호2, 가중치
        distanceData.append((planets[i][3],planets[i+1][3],abs(planets[i][1]-planets[i+1][1])))
    }
    
    // z값 기준 정렬
    planets.sort { (a, b) -> Bool in
        return a[2] < b[2]
    }
    
    for i in 0..<planets.count-1{
        // 행성번호1, 행성번호2, 가중치
        distanceData.append((planets[i][3],planets[i+1][3],abs(planets[i][2]-planets[i+1][2])))
    }
    
    // 추가된 가중치들을 기준으로 정렬
    distanceData.sort { (a, b) -> Bool in
        return a.2 < b.2
    }
    
    // 크루스칼 알고리즘 진행
    for i in 0..<distanceData.count{
        if lines == n-1 {
            break
        }
        if find(distanceData[i].0) != find(distanceData[i].1) {
            lines += 1
            result += distanceData[i].2
            union(distanceData[i].0, distanceData[i].1)
        }
    }
    print(result)
    
}

solution()