스위프트 코딩테스트 강좌, 최소 신장 트리 구하기

오늘은 코딩테스트에서 자주 출제되는 알고리즘 문제 중 하나인 최소 신장 트리(MST: Minimum Spanning Tree) 구하기에 대해 자세히 살펴보겠습니다. 최소 신장 트리는 주어진 가중치가 있는 그래프의 모든 정점을 포함하면서, 가중치의 합이 최소가 되는 트리입니다. 이 문제를 해결하기 위해 우리는 크루스칼 알고리즘(Kruskal’s Algorithm)과 프림 알고리즘(Prim’s Algorithm)을 사용할 수 있습니다. 이 글에서는 크루스칼 알고리즘을 중심으로 문제를 다루도록 하겠습니다.

문제 설명

문제: 주어진 그래프의 최소 신장 트리를 구하세요.

입력으로 주어지는 그래프는 다음과 같은 형식으로 주어집니다:

8 11
0 1 4
0 7 8
1 2 8
1 7 11
2 3 7
2 5 4
2 6 2
3 4 9
3 5 14
4 5 10
5 6 2

첫 줄은 정점의 수와 간선의 수를 나타내며, 두 번째 줄부터는 간선의 정보를 각각 제공합니다. 여기서 각 간선은 두 정점과 가중치로 구성되어 있습니다. 예를 들어, “0 1 4″는 정점 0과 정점 1을 연결하는 간선의 가중치가 4임을 의미합니다.

문제 해결 전략

이번 문제를 해결하기 위해서는 다음과 같은 단계가 필요합니다:

1. 그래프의 간선 정보를 저장합니다.
2. 크루스칼 알고리즘을 적용하여 최소 신장 트리를 구합니다.
3. 최종적으로 선택된 간선들의 가중치를 합산하여 결과를 출력합니다.

1. 그래프의 간선 정보 저장

먼저, 입력을 읽어서 그래프의 간선 정보를 저장합니다. 이를 위해서 각 간선을 (가중치, 시작 정점, 끝 정점) 형태의 튜플로 저장합니다. 각 간선은 가중치에 따라 정렬할 수 있도록 설계해야 합니다.

import Foundation

struct Edge {
    let weight: Int
    let start: Int
    let end: Int
}

func readGraph() -> [Edge] {
    let firstLine = readLine()!.split(separator: " ").map { Int($0)! }
    let vertexCount = firstLine[0]
    let edgeCount = firstLine[1]
    var edges = [Edge]()

    for _ in 0..

2. 크루스칼 알고리즘 구현

크루스칼 알고리즘은 모든 간선을 가중치에 따라 정렬한 후, 가장 낮은 가중치부터 차례로 선택하여 사이클을 구성하지 않도록 추가하는 방식입니다. 이를 위해서 유니온 파인드(Union-Find) 자료구조를 사용합니다. 유니온 파인드는 두 개의 집합이 연결되어 있는지를 빠르게 검사하고, 두 집합을 합치는 기능을 제공합니다.

class DisjointSet {
    private var parent: [Int]

    init(size: Int) {
        self.parent = Array(0.. Int {
        if parent[node] != node {
            parent[node] = find(parent[node])
        }
        return parent[node]
    }

    func union(_ a: Int, _ b: Int) {
        let rootA = find(a)
        let rootB = find(b)
        if rootA != rootB {
            parent[rootB] = rootA
        }
    }
}

func kruskal(edges: [Edge], vertexCount: Int) -> Int {
    let sortedEdges = edges.sorted { $0.weight < $1.weight }
    let disjointSet = DisjointSet(size: vertexCount)
    var totalWeight = 0

    for edge in sortedEdges {
        if disjointSet.find(edge.start) != disjointSet.find(edge.end) {
            disjointSet.union(edge.start, edge.end)
            totalWeight += edge.weight
        }
    }

    return totalWeight
}

3. 메인 함수

이제 모든 부분이 준비되었습니다. 마지막으로 메인 함수를 작성하여 전체 프로세스를 실행합니다.

func main() {
    let edges = readGraph()
    let totalWeight = kruskal(edges: edges, vertexCount: 8)
    print("최소 신장 트리의 가중치 합: \(totalWeight)")
}

main()

전체 프로그램

위의 모든 단계가 완료되었으니, 이제 우리는 전체 프로그램을 다음과 같이 작성할 수 있습니다:

import Foundation

struct Edge {
    let weight: Int
    let start: Int
    let end: Int
}

class DisjointSet {
    private var parent: [Int]

    init(size: Int) {
        self.parent = Array(0.. Int {
        if parent[node] != node {
            parent[node] = find(parent[node])
        }
        return parent[node]
    }

    func union(_ a: Int, _ b: Int) {
        let rootA = find(a)
        let rootB = find(b)
        if rootA != rootB {
            parent[rootB] = rootA
        }
    }
}

func readGraph() -> [Edge] {
    let firstLine = readLine()!.split(separator: " ").map { Int($0)! }
    let vertexCount = firstLine[0]
    let edgeCount = firstLine[1]
    var edges = [Edge]()

    for _ in 0.. Int {
    let sortedEdges = edges.sorted { $0.weight < $1.weight }
    let disjointSet = DisjointSet(size: vertexCount)
    var totalWeight = 0

    for edge in sortedEdges {
        if disjointSet.find(edge.start) != disjointSet.find(edge.end) {
            disjointSet.union(edge.start, edge.end)
            totalWeight += edge.weight
        }
    }
    return totalWeight
}

func main() {
    let edges = readGraph()
    let totalWeight = kruskal(edges: edges, vertexCount: 8)
    print("최소 신장 트리의 가중치 합: \(totalWeight)")
}

main()

결론

이번 포스팅에서는 스위프트를 사용하여 최소 신장 트리 문제를 해결하는 방법을 살펴보았습니다. 크루스칼 알고리즘을 통해 간선 기반으로 MST를 구하는 과정과 유니온 파인드를 사용하는 방법을 알아보았습니다. 실제 코딩 테스트와 면접에서 이러한 문제를 마주쳤을 때 이론과 알고리즘을 바탕으로 자신감을 가지고 접근할 수 있기를 바랍니다. 다음번에는 프림 알고리즘을 이용한 최소 신장 트리 구하기에 대해 다룰 예정이니 많은 기대 부탁드립니다.