2887번: 행성 터널

문제 풀이 아이디어

크루스칼 알고리즘

행성과 행성 사이를 N - 1개의 간선 만을 사용해서 최소 비용으로 연결해야 합니다. 최소 신장 트리를 만드는 크루스칼 알고리즘을 사용해야 하는 것을 알 수 있습니다.

간선을 어떻게 구할 것인가?

모든 간선을 구하면 O(N^2)

일반적인 크루스칼 알고리즘은 보통 간선이 비용과 함께 주어집니다. 하지만 이 문제의 경우 간선의 비용을 구하는 공식만이 주어질 뿐 간선이 직접적으로 주어지지 않습니다.

원칙적으로는 모든 행성 간의 간선을 구해서 비용을 내림차순으로 정렬해서 최소 신장 트리를 구해야만 합니다만 문제의 N은 최대 100,000입니다. 모든 행성 간의 간선을 구한다면 N(N - 1) / 2 개의 간선이 나올 것입니다. 시간 복잡도가 O(N^2)으로 무조건 시간초과 입니다.

3(N - 1)개의 간선만 가지고 크루스칼 알고리즘

다행히도 간선의 비용을 구하는 공식을 보면 간선의 갯수를 획기적으로 줄일 수 있습니다. 간선의 비용은 x, y, z 좌표를 모두 사용해서 구하는 것이 아니라 x, y, z 좌표의 차이 중에서 가장 작은 값을 비용으로 채택하도록 되어 있습니다.

그렇다면 행성을 x 축에 나열된 순서대로 정렬하고 서로 인접한 행성 간의 간선만을 구합니다. N개의 행성이 있다면 N - 1개의 간선이 생길 것입니다. 이 간선만 가지고도 최소 신장트리를 만들 수 있습니다. 추가적으로 y와 z 축을 기준으로도 나열해서 서로 인접한 행성 간의 간선을 구합니다. 총 3(N - 1)개의 간선이 만들어집니다.

이 간선들을 각 축을 기준으로 최소 비용을 가진 간선들만 구한 것입니다. 이 간선들만 가지고 크루스칼 알고리즘을 실행하면 최소 비용을 얻을 수 있습니다.

코드

// 행성과 간선을 튜플로 표현
typealias Planet = (Int, Int, Int, Int)
typealias Edge = (Int, Int, Int)

// 두 행성 간의 간선을 구하는 함수
func getEdges(_ p1: Planet, _ p2: Planet) -> Edge {
    let xCost = abs(p1.1 - p2.1)
    let yCost = abs(p1.2 - p2.2)
    let zCost = abs(p1.3 - p2.3)
    let cost = min(xCost, yCost, zCost)
    return (p1.0, p2.0, cost)
}

// 행성 입력 받기
let N = Int(readLine()!)!
var planets = [Planet]()

for i in 0..<N {
    let xyz = readLine()!.split(separator: " ").map { Int(String($0))! }
    let x = xyz[0], y = xyz[1], z = xyz[2]
    planets.append((i, x, y, z))
}

// 간선 구하기
var edges = [Edge]()

// x 좌표 기준으로 인접한 두 행성의 간선 구하기
planets.sort(by: { $0.1 < $1.1 })
for i in 0..<(N - 1) {
    edges.append(getEdges(planets[i], planets[i + 1]))
}

// y 좌표 기준으로 인접한 두 행성의 간선 구하기
planets.sort(by: { $0.2 < $1.2 })
for i in 0..<(N - 1) {
    edges.append(getEdges(planets[i], planets[i + 1]))
}

// z 좌표 기준으로 인접한 두 행성의 간선 구하기
planets.sort(by: { $0.3 < $1.3 })
for i in 0..<(N - 1) {
    edges.append(getEdges(planets[i], planets[i + 1]))
}

// 3(N - 1)개의 간선들을 비용 순으로 정렬한다.
edges.sort(by: { $0.2 < $1.2 })

// 서로소 집합 알고리즘 구현
var parent = Array(0..<N)

func find(_ a: Int) -> Int {
    if parent[a] != a {
        parent[a] = find(parent[a])
    }
    return parent[a]
}

func union(_ a: Int, _ b: Int) {
    let a = find(a)
    let b = find(b)
    
    if a < b {
        parent[b] = a
    } else {
        parent[a] = b
    }
}

// 크루스칼 알고리즘으로 최소 비용 구하기
var cost = 0

for edge in edges {
    let p1 = edge.0, p2 = edge.1, c = edge.2
    
    if find(p1) != find(p2) {
        union(p1, p2)
        cost += c
    }
}

print(cost)