[Algorithm] 프로그래머스 - 배달 (Swift) (플로이드 워셜, DFS)

문제 소개

프로그래머스

N개의 마을이 있고, 이 마을간에는 도로가 연결되어있다. 그리고 각 마을간 최소거리가 k이하인 마을의 갯수를 구하는 문제이다.

 

출처: 프로그래머스

플로이드-워셜이나 다익스트라 알고리즘을 사용하면 아주 손쉽게 풀 수 있다. 나는 플로이드-워셜 알고리즘을 사용했다. 하지만 이것은 너무나도 치트키 같은 느낌... 특히 문제를 풀어본뒤 질문들을 보는 습관이 생겼는데 BFS, DFS로도 충분히 풀수있는 사이즈이기 때문에 도전해보고자 했다.

 

플로이드 워셜 알고리즘

아주아주 코드가 쉽다. 3중 for문만 시원하게 써주면 문제풀이 끝! 주의할점은 처음 table을 만들어줄때 초기값을 잘 설정해야 한다는 것? 무한으로 설정하기도 하는데 Swift의 경우에는 Int.max로 선언이 가능하고 922,3372,0368,5477,5807 가 최댓값이다.

 

하지만 그렇게되면 int로 설정할 수 있는 최댓값을 넘을 수 있는 가능성이 있기 때문에 아래 핸들링을 해주거나 초기값을 문제를 보고 절대 나올수 없는 최댓값으로 해주면 될 듯 하다. (개인적으로 두번째 방법을 선호한다... 현실세계에서 저 숫자까지 올릴 필요가 있을까...?) 그리고 내자신으로 가는것은 0으로 설정해주고 3중 for문을 해주면 된다. 

func solution(_ N:Int, _ road:[[Int]], _ k:Int) -> Int {
    var answer = 0
    var table = Array(repeating: Array(repeating: 500000, count: N), count: N)
    
    for r in road {
        let i = r[0]-1
        let j = r[1]-1
        
        if table[i][j] > r[2] {
            table[i][j] = r[2]
            table[j][i] = r[2]
        }
    }
    
    for i in 0..<N {
        table[i][i] = 0
    }
        
    for k in 0..<N {
        for i in 0..<N {
            for j in 0..<N {
                table[i][j] = min(table[i][j], table[i][k]+table[k][j])
            }
        }
    }
    
    for i in 0..<N {
        if table[0][i] <= k {
            answer += 1
        }
    }

    return answer
}

 

DFS(깊이 우선 탐색)

문제에서 마을의 갯수가 최대 50개라는 점, 도로정보도 최대 2000개라는 점...을 고려하면 아무리 2중 for문을 돌려도 최악의 경우 100000번 탐색을 하게 되는데 몇가지 조건만 잘 설정해준다면 DFS도 충분히 가능할 것이라고 생각했다.

 

첫 풀이

아래와 같이 cost를 더해가면서 만약 cost가 k보다 크다면 종료, 이미 간 목표지점이면 종료, 이것을 제외한 경우 목표지점에 도달하면 정답에 하나를 더해주는 식으로 dfs함수를 구성했다. 

func solution(_ N:Int, _ road:[[Int]], _ k:Int) -> Int {
    var answer = 0
    var ends = [Int: Bool]()
    
    func dfs(_ cost: Int, _ s: Int, _ e: Int, _ v: [Int: Bool]) {
        if cost > k { return }
        if ends[e] != nil { return }
        if s == e {
            ends[e] = true
            answer += 1
            return
        }

        for r in road {
            var v = v
            var newS = 0
            var newCost = 10001
            
            if r[2] < newCost {
                if r[0] == s {
                    if v[r[1]] == nil {
                        newS = r[1]
                        newCost = r[2]
                    }
                } else if r[1] == s {
                    if v[r[0]] == nil {
                        newS = r[0]
                        newCost = r[2]
                    }
                }
            }
            
            if newCost < 10001 {
                v[newS] = true
                dfs(cost+newCost, newS, e, v)
            }
        }
    }
    
    for i in 1...N {
        dfs(0, 1, i, [Int: Bool]())
    }
    
    return answer
}

 

실행 결과 및 수정

실행 결과 테스트 1~11까지는 0.1ms 도 안걸릴 정도로 짧았지만, 문제는 테스트케이스 32번이 시간 초과가 되었다.

 

내가 생각하기에 문제는 마지막 부분에 있었다. dfs자체를 N번 반복하게 되니까... 비효율적이라고 생각했다. 어자피 시작노드는 정해져있고, 모든 노드를 방문할때까지 순회하면 되는 것이라고 생각했다. 

 

다시 접근했다. 우선 nodes를 만들어 key값으로 접근할 수 있게 하였고, ends에도 key값과 드는 시간을 저장했다. 여기서 중요한 것은 함수의 탈출조건인데, 우선 드는 시간이 최대시간보다 크면 탈출시키고, ends에 저장되어있는 시간 값이랑 비교해서 cost가 크면 바로 return 시켜버린다. 왜냐하면 k보다는 작지만 cost가 최신 업데이트된 end값보다 크면 의미가 없기 때문이다.

import Foundation

func solution(_ N:Int, _ road:[[Int]], _ k:Int) -> Int {
    var ends = [Int: Int]()
    var nodes = [Int: [(node: Int, cost: Int)]]()
    
    for i in 1...N {
        nodes[i] = []
    }
    
    for r in road {
        nodes[r[0]]?.append((node: r[1], cost: r[2]))
        nodes[r[1]]?.append((node: r[0], cost: r[2]))
    }
        
    func dfs(_ cost: Int, _ current: Int) {
        if cost > k { return }
        
        if ends[current] == nil {
            ends[current] = cost
        } else if let end = ends[current] {
            ends[current] = min(end, cost)
            if end < cost {
                return
            }
        }

        for next in nodes[current]! {
            let nextNode = next.node
            var nextCost = next.cost
            
            dfs(cost+nextCost, nextNode)
        }
    }
    
    dfs(0, 1)
    
    return ends.count
}

 

실행결과는 아래와 같이 훨~씬 개선된 결과가 나왔다 :)