[백준] 17472번: 다리 만들기 2

문제링크

• DFS, BFS 탐색
• MST 알고리즘

import java.util.*

data class Edge(val v1: Int, val v2: Int, val w: Int)

lateinit var world: MutableList<MutableList<Int>>
lateinit var edges: PriorityQueue<Edge>
lateinit var parent: Array<Int>

val dx = listOf(1, 0, -1, 0)
val dy = listOf(0, 1, 0, -1)

fun main() {
    val br = System.`in`.bufferedReader()
    val bw = System.out.bufferedWriter()

    world = mutableListOf()
    val (n, m) = br.readLine().split(" ").map { it.toInt() }
    for (i in 0 until n) {
        world.add(br.readLine().split(" ").map { it.toInt() }.toMutableList())
    }

    // 섬을 구별하는 숫자가 1이므로 2부터 시작합니다.
    var island = 2
    for (i in 0 until n) {
        for (j in 0 until m) {
            if (world[i][j] == 1) {
                // 섬이면, BFS 탐색을 통해 인접한 섬을 동일한 숫자로 표시합니다.
                bfs(n, m, i, j, island)
                island++
            }
        }
    }
    
    // 섬 간의 거리를 오름차순으로 저장합니다.
    edges = PriorityQueue(compareBy { edge -> edge.w })
    for (i in 0 until n) {
        for (j in 0 until m) {
            // 섬이면, 다른 섬과의 거리를 계산하여 우선순위 큐에 저장합니다.
            if (world[i][j] != 0) findEdge(n, m, i, j)
        }
    }

    // MST (kruskal)
    parent = Array(island) { it }
    var answer = 0
    var cnt = 0
    while (edges.isNotEmpty()) {
        val edge = edges.poll()
        val p1 = find(edge.v1)
        val p2 = find(edge.v2)
        if (p1 != p2) {
            answer += edge.w
            union(p1, p2)
            cnt++
        }
    }

    // island starts with 2 and cnt must be equal to |V| - 1
    // if island -> 5, means |V| -> 3
    // so, cnt must be 2. Thus, cnt == island - 2
    if (cnt == island - 3) {
        bw.write("$answer")
    } else {
        bw.write("-1")
    }

    br.close()
    bw.close()
}

fun bfs(n: Int, m: Int, x: Int, y: Int, island: Int) {
    val visited = Array(n) { Array(m) { false } }
    val q: Queue<Pair<Int, Int>> = LinkedList()
    q.add(Pair(x, y))
    world[x][y] = island
    visited[x][y] = true

    while (q.isNotEmpty()) {
        val (cx, cy) = q.poll()
        for (i in 0..3) {
            val nx = cx + dx[i]
            val ny = cy + dy[i]
            if (nx !in 0 until n || ny !in 0 until m) continue
            if (visited[nx][ny] || world[nx][ny] != 1) continue

            visited[nx][ny] = true
            q.add(Pair(nx, ny))
            world[nx][ny] = island
        }
    }
}

fun findEdge(n: Int, m: Int, x: Int, y: Int) {
    // 위, 아래, 왼쪽, 오른쪽 방향으로 직진합니다.
    for (i in 0..3) {
        var nx = x
        var ny = y
        var w = 0
        while (true) {
            nx += dx[i]
            ny += dy[i]
            // 범위 밖 좌표 무시
            if (nx !in 0 until n || ny !in 0 until m) break
            // 같은 섬 무시
            if (world[nx][ny] == world[x][y]) break
            // 바다일 때에만 거리 증가
            if (world[nx][ny] == 0) w++
            // 거리가 1 이라면 무시
            else if (w == 1) break
            else {
                // 다른 섬 도착 -> 현재 섬과 도착 섬 정보와 거리 저장
                edges.add(Edge(world[x][y], world[nx][ny], w))
                break
            }
        }
    }
}

fun find(x: Int): Int {
    if (x == parent[x]) return x
    parent[x] = find(parent[x])
    return parent[x]
}

fun union(x: Int, y: Int) {
    val px = find(x)
    val py = find(y)
    if (px != py) {
        parent[py] = px
    }
}

시간복잡도

O(nm)