코딩테스트 연습 - 배달

문제 풀이 아이디어

플로이드 와샬 알고리즘

문제를 풀기 위해서는 일단 1번 마을과 다른 마을 사이의 최단경로를 구해야 합니다. 1번 마을과 직접 연결된 도로 뿐만이 아니라 다른 마을을 건너서 가는 경로가 최단 경로가 될 수도 있습니다. 따라서 모든 node 사이의 최단 경로를 찾아주는 플로이드 와샬 알고리즘을 사용해서 풀도록 하겠습니다. 플로이드 와샬 알고리즘에 대한 설명은 이 포스팅을 참고해주세요.

이차원 배열의 초기값 설정에 주의할 점

문제의 조건을 보면 도로의 최대 길이는 10000입니다. 따라서 도로가 없음을 표현하기 위해서 10001을 이차원 배열의 초기값으로 사용할 수도 있는데요. 이 경우 K의 값이 큰 경우 코드가 통과되지 않습니다. K는 최대 500,000입니다. 만약 K의 값이 최댓값인 500,000인 경우 도로가 없음을 10001로 표현하는 경우 실제로는 도로가 없음에도 배달이 가능한 경로가 생기게 됩니다. 따라서 K의 최댓값을 고려해서 이차원 배열의 초기값(= 도로 없음)은 500,001로 하는 것이 좋습니다.

코드

func solution(_ N:Int, _ road:[[Int]], _ k:Int) -> Int {
    // 이차원 배열 세팅
    var matrix = Array(repeating: Array(repeating: 500001, count: N + 1), count: N + 1)

    // 자기마을까지의 경로의 비용 0으로 세팅
    for i in 1...N {
        matrix[i][i] = 0
    }

    // 주어진 양방향 도로 이차원 배열에 입력
    for r in road {
        let a = r[0]
        let b = r[1]
        let c = r[2]
        matrix[a][b] = min(matrix[a][b], c)
        matrix[b][a] = min(matrix[b][a], c)
    }

    // 플로이드 와샬 알고리즘
    for k in 1...N {
        for i in 1...N {
            for j in 1...N {
                matrix[i][j] = min(matrix[i][j], matrix[i][k] + matrix[k][j])
            }
        }
    }

    // 1에서 다른 마을까지 k보다 짧은 경우 세기
    var ans = 0

    for i in 1...N {
        if matrix[1][i] <= k { ans += 1 }
    }

    return ans
}