[프로그래머스 kotlin] 합승 택시 요금

문제

 

시작지점 S에서 출발하여 2명의 사람이 각각 A지점과 B지점에 가야할 때 최단거리를 구하는 문제이다.

---

풀이

시작지점에서 출발하여 A, B를 돌도록 최단거리를 구하면 안된다.

그 이유는 한 지점에서 각각의 사람이 갈라져서 이동하기 때문이다.

예를들어 위의 사진의 경우 S에서 출발하여 1번노드로 이동후 B는 5 -> 3 -> 2 노드로 이동하고 A는 5 -> 6 노드로 이동하므로
각각의 경우를 따로따로 구해야 한다.

여기서 봐야할 점은 S에서 출발하여 A, B 로 각각 간다는점인데 이를 다시생각하면 한 지점에서 S, A, B지점으로 이동한 거리를 모두 합한 값이 된다.

 

따라서 각 지점에 대해 S, A, B 까지의 최단거리를 합한 후 가장 작은 값을 구하면 된다.

 

 

data class Node(val end: Int, val weight: Int): Comparable<Node> {
    override fun compareTo(other: Node) = weight - other.weight
}

weight값을 기준으로 정렬하는 data class Node를 정의해준다.

 

private lateinit var dst: IntArray
private lateinit var list: Array<ArrayList<Node>>

 

거리값을 저장한 dst와 간선정보를 저장할 list를 선언해준다.

 

var answer: Int = Int.MAX_VALUE
list = Array(n + 1) { ArrayList<Node>() }
for (i in 0 until fares.size) {
    val (start, end, weight) = fares[i].map { it }
    list[start].add(Node(end, weight))
    list[end].add(Node(start, weight))
}

정답을 출력할 answer는 최대값으로 초기화하고

간선의 정보를 저장할 list는 n + 1의 크기로 초기화한다. 해당 배열은 Node를 가지는 ArrayList를 가진다.

따라서 [ [Node, Node], [Node], [Node, Node, Node] ] 의 형태를 가진다.

 

이후 해당 list에 간선의 정보를 입력받는다. 해단 간선은 양방향이기 때문에 양쪽 방향을 모두 저장한다.

fun dijkstra(start: Int, n: Int): IntArray {
    val q = PriorityQueue<Node>()
    val isVisited = BooleanArray(n + 1)
    dst = IntArray(n + 1) { Int.MAX_VALUE }
    dst[start] = 0
    q.offer(Node(start, 0))

    while (q.isNotEmpty()) {
        val cur = q.poll()
        if (isVisited[cur.end]) continue
        isVisited[cur.end] = true

        for (next in list[cur.end]) {
            if (dst[next.end] > dst[cur.end] + next.weight) {
                dst[next.end] = dst[cur.end] + next.weight
                q.offer(Node(next.end, dst[next.end]))
            }
        }
    }
    return dst
}

다익스트라의 구현부이다. 

다익스트라를 여러번 실행해야하기 때문에 해당 함수 내부에서 visited를 관리해주었다.

dst 값을 초기화해주고 현재 시작위치의 거리를 0으로 해준다.

그 후 priority queue인 q에 넣어주고 while문을 실행한다.

이후는 기존 다익스트라의 구현과 같다.

 

마지막으로 값을 할당한 dst를 return한다.

 

val dstS = dijkstra(s, n)
val dstA = dijkstra(a, n)
val dstB = dijkstra(b, n)

for (i in 1..n) {
    if (dstS[i] == Int.MAX_VALUE || dstA[i] == Int.MAX_VALUE || dstB[i] == Int.MAX_VALUE) continue
    answer = min(answer, dstS[i] + dstA[i] + dstB[i])
}

다익스트라를 사용하여 거리를 구하는 부분이다.

dstS, A, B 를 각각의 지점에서 시작하는 최단거리값을 할당해준다.

 

마지막으로 모든 값을 더했을 때의 최소값을 answer에 할당한다.

---

코드

import java.util.*
import kotlin.math.*

class Solution {
    private lateinit var dst: IntArray
    private lateinit var list: Array<ArrayList<Node>>
    
    fun solution(n: Int, s: Int, a: Int, b: Int, fares: Array<IntArray>): Int {
        var answer: Int = Int.MAX_VALUE
        list = Array(n + 1) { ArrayList<Node>() }
        for (i in 0 until fares.size) {
            val (start, end, weight) = fares[i].map { it }
            list[start].add(Node(end, weight))
            list[end].add(Node(start, weight))
        }
        
        val dstS = dijkstra(s, n)
        val dstA = dijkstra(a, n)
        val dstB = dijkstra(b, n)
        
        for (i in 1..n) {
            if (dstS[i] == Int.MAX_VALUE || dstA[i] == Int.MAX_VALUE || dstB[i] == Int.MAX_VALUE) continue
            answer = min(answer, dstS[i] + dstA[i] + dstB[i])
        }
        
        
        
        return answer
        
    }
    
    fun dijkstra(start: Int, n: Int): IntArray {
        val q = PriorityQueue<Node>()
        val isVisited = BooleanArray(n + 1)
        dst = IntArray(n + 1) { Int.MAX_VALUE }
        dst[start] = 0
        q.offer(Node(start, 0))
        
        while (q.isNotEmpty()) {
            val cur = q.poll()
            if (isVisited[cur.end]) continue
            isVisited[cur.end] = true
            
            for (next in list[cur.end]) {
                if (dst[next.end] > dst[cur.end] + next.weight) {
                    dst[next.end] = dst[cur.end] + next.weight
                    q.offer(Node(next.end, dst[next.end]))
                }
            }
        }
        return dst
    }
}

data class Node(val end: Int, val weight: Int): Comparable<Node> {
    override fun compareTo(other: Node) = weight - other.weight
}