[프로그래머스/Swift] 배달

 

https://school.programmers.co.kr/learn/courses/30/lessons/12978

프로그래머스

 

이번 문제를 보고 처음에는 최단 거리를 구하는 문제니 BFS로 풀려고 했다.

그런데 풀다 보니 전에 C++로 풀었던 다익스트라를 사용해야겠구나! 생각했다. 그런데 priority_queue STL로 없어서 직접 구현해야 했다.. 코테를 풀 때 이래서 c++ 나 파이썬을 사용하는구나... 이걸 언제 다 구현하고 앉아있지..? Swift한테 너무 가혹한 거 아닌가요?ㅠㅠ

 

다익스트라 자체는 생각보다 쉽다! 간단히 정리해보자면

weight가 있는 그래프에서 최단 거리를 구하려고 할 때 사용한다.

→ 현재까지 알고 있던 최단 경로를 계속해서 갱신한다. Heap를 가져야 빠른 시간 내에 풀 수 있다.

dist 배열을 통해 시작 노드부터 각 노드의 거리를 구한다. 이 값을 계속 갱신하는 것이다!!

 

1. 시작 노드를 제외하고 모든 노드의 dist 값을 무한대(inf)로 넣는다. (Int.max)

2. 아직 방문하지 않은 노드 중 가장 비용이 적은 노드를 선택한다. (priority_queue를 사용해서 우선순위(비용 적은 순)를 따른다)

3. 해당 노드를 거쳐서 특정 노드로 가는 경우를 고려해서 min cost 갱신한다.

 

💡 만약 노드의 비용(now.weight)보다 지금까지 구한 현재 노드의 비용(dist [now.vertex])가 크다면 패스! 더 적은 비용을 구해야 하니까! 바로 버려주자~

 

현재 노드의 비용(dist[now.vertex]) + 다음 노드로 가는 비용(dist [next.weight]) < 지금까지 구한 다음 노드로 가는 비용(dist [next.vertex]) 라면 업데이트 시켜준다! 더 적은 비용으로 가야 하니까

 

이렇게 알아보면 쉬운데 이제 우선순위 큐를 구현해야 한다. 우선순위 큐는 Heap으로 구현되므로 Heap 구조체를 작성해 준다.

혼자서 만들어보려다가 처음에는... 다른 분의 코드를 보는 게 좋을 거 같아서 찾아봤다.

[라이노 님의 Heap 코드다]

 

public struct Heap<T> {
    var nodes: [T] = [] //heap nodes 담을 배열
    let comparer: (T,T) -> Bool //비교자
    
    var isEmpty: Bool {
        return nodes.isEmpty
    }
    
    init(comparer: @escaping (T,T) -> Bool) { //비교자 초기화
        self.comparer = comparer
    }
    
    func peek() -> T? {
        return nodes.first
    }
    
    //노드 추가
    mutating func insert(_ element: T) {
        var index = nodes.count //노드가 추가될 맨 뒤 인덱스
        
        nodes.append(element) //노드 추가
        
        //(index-1)/2 == 부모 노드
        //index == 0 (꼭대기 노드) 되거나 부모보다 크(<)/작으(>)면 끝
        while index > 0, !comparer(nodes[index], nodes[(index-1)/2]) {
            nodes.swapAt(index, (index-1)/2)
            index = (index - 1) / 2
        }
    }
    
    mutating func delete() -> T? {
        guard !nodes.isEmpty else { return nil }
        if nodes.count == 1 { return nodes.removeFirst() }
        
        let result = nodes.first //삭제될 맨 첫번째(min/max) element
        nodes.swapAt(0, nodes.count - 1) //맨 위랑 맨 마지막 swap
        _ = nodes.popLast() //last element 삭제
        
        var index = 0
        
        
        //모든 노드 이동..
        while index < nodes.count { //배열 모두 순회
            let left = index * 2 + 1
            let right = left + 1
            
            if right < nodes.count {
                if comparer(nodes[left], nodes[right]),
                   !comparer(nodes[right], nodes[index]) {
                    nodes.swapAt(right, left)
                    index = right
                } else if !comparer(nodes[left], nodes[index]) {
                    nodes.swapAt(left, right)
                    index = left
                } else {
                    break
                }
            } else if left < nodes.count {
                if !comparer(nodes[left], nodes[index]) {
                    nodes.swapAt(left, index)
                    index = left
                } else {
                    break
                }
            } else {
                break
            }
        }
        return result
    }
}

extension Heap where T: Comparable {
    init() {
        self.init(comparer: >) //<: max heap, >: min heap
    }
}

 

자료구조 시간 때 배운 Min/Max Heap이다. 여기서 extension의 comparer의 부호만 바꿔주면 max 혹은 min heap을 구현할 수 있다.

이 Heap을 사용해서 priority queue를 구현하고, vertex와 weight 정보는 struct로 내가 custom 해줬다.

 

 

[전체 코드]

import Foundation

public struct Heap<T> {
    var nodes: [T] = [] //heap nodes 담을 배열
    let comparer: (T,T) -> Bool //비교자
    
    var isEmpty: Bool {
        return nodes.isEmpty
    }
    
    init(comparer: @escaping (T,T) -> Bool) { //비교자 초기화
        self.comparer = comparer
    }
    
    //노드 추가
    mutating func insert(_ element: T) {
        var index = nodes.count //노드가 추가될 맨 뒤 인덱스
        
        nodes.append(element) //노드 추가
        
        //(index-1)/2 == 부모 노드
        //index == 0 (꼭대기 노드) 되거나 부모보다 크(<)/작으(>)면 끝
        while index > 0, !comparer(nodes[index], nodes[(index-1)/2]) {
            nodes.swapAt(index, (index-1)/2)
            index = (index - 1) / 2
        }
    }
    
    mutating func delete() -> T? {
        guard !nodes.isEmpty else { return nil }
        if nodes.count == 1 { return nodes.removeFirst() }
        
        let result = nodes.first //삭제될 맨 첫번째(min/max) element
        nodes.swapAt(0, nodes.count - 1) //맨 위랑 맨 마지막 swap
        _ = nodes.popLast() //last element 삭제
        
        var index = 0
        
        
        //모든 노드 이동..
        while index < nodes.count { //배열 모두 순회
            let left = index * 2 + 1
            let right = left + 1
            
            if right < nodes.count {
                if comparer(nodes[left], nodes[right]),
                   !comparer(nodes[right], nodes[index]) {
                    nodes.swapAt(right, left)
                    index = right
                } else if !comparer(nodes[left], nodes[index]) {
                    nodes.swapAt(left, right)
                    index = left
                } else {
                    break
                }
            } else if left < nodes.count {
                if !comparer(nodes[left], nodes[index]) {
                    nodes.swapAt(left, index)
                    index = left
                } else {
                    break
                }
            } else {
                break
            }
        }
        return result
    }
}

extension Heap where T: Comparable {
    init() {
        self.init(comparer: >) //<: max heap, >: min heap
    }
}

struct Edge: Comparable {
    static func < (lhs: Edge, rhs: Edge) -> Bool {
        lhs.weight < rhs.weight
    }
    
    let vertex: Int
    let weight: Int
}

struct Dict : Hashable {
    let vertex: Int
    let weight: Int
}

func solution(_ N:Int, _ road:[[Int]], _ k:Int) -> Int {
    var answer = 0
    let inf = Int.max
    var dist = Array(repeating: inf, count: N+1)
    var graph = [[Dict]](repeating: [Dict](), count: N + 1)
    road.forEach {
        graph[$0[0]].append(Dict(vertex: $0[1], weight: $0[2]))
        graph[$0[1]].append(Dict(vertex: $0[0], weight: $0[2]))
    }
    
    var pq : Heap = Heap<Edge>()
    pq.insert(Edge(vertex: 1, weight: 0))
    dist[1] = 0 //시작점 0으로 초기화

    while !pq.isEmpty {
        let now = pq.delete()! //weight가 가장 작은 애 없앰
        if dist[now.vertex] < now.weight { continue }

        for next in graph[now.vertex] {
            if now.weight + next.weight < dist[next.vertex] {
                dist[next.vertex] = now.weight + next.weight
                pq.insert(Edge(vertex: next.vertex, weight: next.weight + now.weight))
            }
        }
    }
    
    for i in 1...N {
        if dist[i] <= k { answer += 1 }
    }
    return answer
}

 

진짜 너무 길고... 어지러울 수 있는 코드다. 하지만 차근차근해보자.. 할 수 있다..

사실 다익스트라 부분은 굉장히 짧기 때문에... 원리만 이해하면 잘할 수 있다!

Heap 부분.. 코테에 나오면 잘 구현할 수 있을까.. c++로 돌려서 해야 하나ㅋㅋㅋ

 

 

---

 

 

이 글을 보시는 분들... 저의 코드는 잊어주세요..

PQ를 그냥 배열로 해도 되는군요... 매번 정렬해 주면서 진행하는 것도 되네요. 아마 제한 범위가 작아서 그런 듯..?

세상에... 이렇게 하나 또 배워갑니다..^^

 

그래도 heap과 다익스트라 공부도 해보고 좋네요

 

import Foundation

var INF = 2100000000

func dijkstra(_ node: Int, _ G: inout [[(Int, Int)]], _ D: inout [Int]) {
    var PQ = [(Int,Int)]()
    PQ.append((1,D[node]))

    while(!PQ.isEmpty) {
        PQ.sort { $0.1 > $1.1 }
        let nowNode = PQ.last!.0
        let nowDistance = PQ.last!.1
        PQ.removeLast()
        for i in G[nowNode].indices {
            let nextNode = G[nowNode][i].0
            let nextDistance = G[nowNode][i].1 + nowDistance
            if nextDistance < D[nextNode] {
                D[nextNode] = nextDistance
                PQ.append((nextNode, nextDistance))
            }
        }
    }
}

func solution(_ N:Int, _ road:[[Int]], _ k:Int) -> Int {
    var answer = 0
    var G = [[(Int, Int)]](repeating: [], count: N+1)
    var distance = [Int](repeating: INF, count: N+1)
    distance[1] = 0

    for i in road.indices {
        G[road[i][0]].append((road[i][1],road[i][2]));
        G[road[i][1]].append((road[i][0],road[i][2]));
    }

    dijkstra(1, &G, &distance)
    for i in 1...N { 
        if distance[i] <= k { answer += 1 }
    }

    return answer
}

 

깔끔하고... 보기좋은 코드...