[BOJ / Swift & Python] 23354 - 군탈체포조

* 개인 블로그에 작성한 내용을 복사해왔습니다. 

[Algorithm] BOJ 23354

BOJ 23354번 군탈체포조

23354번: 군탈체포조

Intro

다익스트라를 사용하여 풀이하면 되는 문제입니다. 처음에는 순열을 사용하여 각 순서대로 다익스트라를 돌려 풀이하였는데, 당연하게도 순열 결과 배열의 길이 * 탈영병 인원수만큼 다익스트라를 돌리다 보니 시간 초과가 떴습니다.
그래서 탈영병 인원만큼만 다익스트라를 돌려 각 탈영병들 간의 최소 비용을 미리 저장해둔 뒤, 순열을 사용하여 전체 최소 비용을 구할 때는 저장해둔 값을 불러와 계산하도록 변경하여 시간 초과를 해결하였습니다.

Solution

1. 모든 입력을 저장한 뒤, 지도에서 탈영병과 부대 위치를 찾아 저장한다. 부대 위치는 0으로 변경시킨다.
2. 탈영병이 없다면 0을 출력하고 종료한다.
3. 부대와 탈영병 그리고 탈영병과 탈영병 간의 거리를 저장하는 (탈영병 인원+1)*(탈영병 인원+1) 크기의 이중 배열을 생성한다. (0번째 인덱스는 부대라고 가정한다.)
4. 각 탈영병들을 시작점으로 다익스트라를 돌려 부대와 탈영병, 그리고 탈영병과 탈영병 간의 거리를 저장한다.
5. 1, 2, ..., 탈영병_인원수가 들어있는 배열의 가능한 순열을 계산한다.
6. 순열의 결과 중 최소 비용이 드는 순서를 찾기 위해 반복문을 돌린다.
   • 탈영병 1명
     부대 -> 첫 번째 탈영병 + 첫 번째 탈영병 -> 부대 의 비용
   • 탈영병 N명( N>1 )
     부대 -> 첫 번째 탈영병 + 첫 번째 탈영병 -> 두 번째 탈영병 + ... + N-1 번째 탈영병 -> N 번째 탈영병 + N 번째 탈영병 -> 부대의 비용
7. 반복문 결과 계산된 최소 비용의 값을 출력한다.

Code

Swift

*순열, 최소 힙, 다익스트라를 모두 구현하여 코드가 조금 깁니다.

import Foundation

struct Coordinate {
    let x: Int
    let y: Int
    func movableLocation(_ N: Int) -> [Self] {
        return [Coordinate(x: self.x-1, y: self.y), Coordinate(x: self.x+1, y: self.y), Coordinate(x: self.x, y: self.y-1), Coordinate(x: self.x, y: self.y+1)].filter() {
            0..<N ~= $0.x && 0..<N ~= $0.y
        }
    }
} 

struct MinHeap<T: Comparable> {
    private var elements: [T] = []
    init(elements: [T]) {
        for e in elements { self.insert(e) }
    }
    var count: Int {
        return self.elements.count
    }
    private var lastIndex: Int {
        return self.count - 1
    }
    private func rightChild(of index: Int) -> Int? {
        guard index*2+2 < self.count else { return nil }
        return index*2+2
    }
    private func leftChild(of index: Int) -> Int? {
        guard index*2+1 < self.count else { return nil }
        return index*2+1
    }
    private func parent(of index: Int) -> Int? {
        guard index > 0 else { return nil }
        return (index-1)/2
    }
    private mutating func swap(_ index1: Int, _ index2: Int) {
        let temp = self.elements[index1]
        self.elements[index1] = self.elements[index2]
        self.elements[index2] = temp
    }
    func node(of index: Int) -> T {
        return self.elements[index]
    }
    mutating func insert(_ node: T) {
        self.elements.append(node)
        var child = self.lastIndex
        while let parents = parent(of: child),
            self.elements[parents] > node {
            swap(parents, child)
            child = parents
        }
    }
    mutating func delete() {
        swap(self.lastIndex, 0)
        self.elements.remove(at: self.lastIndex)
        var current = 0
        var lChild = leftChild(of: current)
        var rChild = rightChild(of: current)
        while current < self.count, 
        self.elements[current] > self.elements[lChild ?? current] || self.elements[current] > self.elements[rChild ?? current] {
            if let lChild = lChild,
            let rChild = rChild {
                if self.elements[lChild] < self.elements[rChild] {
                    swap(current, lChild)
                    current = lChild
                } else {
                    swap(current, rChild)
                    current = rChild
                }
            } else if let lChild = lChild {
                swap(current, lChild)
                current = lChild
            } else if let rChild = rChild {
                swap(current, rChild)
                current = rChild
            } else { break }
            lChild = leftChild(of: current)
            rChild = rightChild(of: current)
        }
    }
    mutating func pop() -> T {
        let removeElement = self.elements[0]
        self.delete()
        return removeElement
    }
}

struct Node: Comparable {
    var element: Int
    var coordinate: Coordinate
    static func < (lhs: Self, rhs: Self) -> Bool {
        return lhs.element < rhs.element
    }
    static func == (lhs: Self, rhs: Self) -> Bool {
        return lhs.element == rhs.element && lhs.coordinate.x == rhs.coordinate.x && lhs.coordinate.y == rhs.coordinate.y
    }
}

func dijkstra(_ N: Int, s: Coordinate, graph: [[Int]]) -> [[Int]] {
    var queue = MinHeap<Node>(elements: [Node(element: graph[s.x][s.y], coordinate: s)])
    var dists = Array<[Int]>(repeating: Array<Int>(repeating: Int.max, count: N), count: N)
    while queue.count > 0 {
        let c = queue.pop()
        let curDist = c.element
        let curLoc = c.coordinate
        guard dists[curLoc.x][curLoc.y] >= curDist else { continue }
        for nextLoc in curLoc.movableLocation(N) {
            let dist = curDist + graph[nextLoc.x][nextLoc.y]
            if dist < dists[nextLoc.x][nextLoc.y] {
                dists[nextLoc.x][nextLoc.y] = dist
                queue.insert(Node(element: dist, coordinate: nextLoc))
            }
        }
    }
    return dists
} 

func permutation<T>(_ elements: [T], _ n : Int) -> [[T]] {
    guard n > 0 else { return [] }
    var result = [[T]]()
    var visited = [Bool](repeating: false, count: elements.count)
    
    func permute(_ now: [T]) {
        if now.count == n {
            result.append(now)
            return
        }
        for i in 0..<elements.count where !visited[i] {
            visited[i] = true
            permute(now+[elements[i]])
            visited[i] = false
        }
    }
    
    permute([])
    return result
}

func minDist(_ N: Int, start: Coordinate, graph: [[Int]], deserters: [Coordinate]) -> Int {
    guard !deserters.isEmpty else { return 0 }
    var minDist = Int.max
    var dists = Array<[Int]>(repeating: Array<Int>(repeating: Int.max, count: deserters.count+1), count: deserters.count+1)
    dists[0][0] = 0
    for i in 1...deserters.count {
        let idist = dijkstra(N, s: deserters[i-1], graph: graph)
        dists[i][0] = idist[start.x][start.y]
        dists[0][i] = dists[i][0]
        for j in 0..<deserters.count {
            dists[i][j+1] = idist[deserters[j].x][deserters[j].y]
            dists[j+1][i] = dists[i][j+1]
        }
    }
    for order in permutation(Array(1...deserters.count), deserters.count) {
        var dist = 0
        var current = 0
        for dest in order {
            dist += dists[current][dest]
            current = dest
        }
        dist += dists[current][0]
        minDist = min(dist, minDist)
    }
    return minDist
}

let N = Int(readLine()!)!
var graph = Array<[Int]>()
var unit: Coordinate? = nil
var deserters: [Coordinate] = []
for i in 0..<N {
    graph.append(readLine()!.split(separator: " ").map() { Int(String($0))! })
    for j in 0..<N {
        switch graph[i][j] {
        case -1:
            graph[i][j] = 0
            unit = Coordinate(x: i, y: j)
        case 0:
            deserters.append(Coordinate(x: i, y: j))
        default: continue
        }
    }
}
print(minDist(N, start: unit!, graph: graph, deserters: deserters))

Python

* 한 게시글에 코드 블록 언어를 하나로 통일해야 해서 Python 코드는 하이라이트가 정상적이지 않습니다.

*시간 초과로 PyPy3로 제출하여 통과하였습니다.

import heapq, sys
from itertools import permutations as pm
input = sys.stdin.readline

n = int(input().strip())

def dijkstra(s, graph) :
    dists = [[float('inf')] * n for _ in range(n)]
    dists[s[0]][s[1]] = 0
    queue = [[dists[s[0]][s[1]], s[0], s[1]]]
    heapq.heapify(queue)
    while queue:
        c = heapq.heappop(queue)
        c_dist = c[0]
        c_dest = (c[1], c[2])
        if dists[c_dest[0]][c_dest[1]] < c_dist : continue
        for n_dest_x, n_dest_y in [(c_dest[0]-1, c_dest[1]), (c_dest[0]+1, c_dest[1]), (c_dest[0], c_dest[1]-1), (c_dest[0], c_dest[1]+1)] :
            if not (0 <= n_dest_x < n and 0 <= n_dest_y < n) : continue
            dist = c_dist + graph[n_dest_x][n_dest_y]
            if dist < dists[n_dest_x][n_dest_y] :
                dists[n_dest_x][n_dest_y] = dist
                heapq.heappush(queue, [dist, n_dest_x, n_dest_y])
    return dists

def minDist(unit, graph, deserters) :
    if not deserters : return 0
    minDist = float('inf')
    dists = [[float('inf')] * (len(deserters)+1) for _ in range(len(deserters)+1)]
    dists[0][0] = 0
    for i in range(1, len(deserters)+1) :
        idist = dijkstra(deserters[i-1], graph)
        for j in range(len(deserters)+1) :
            if j == 0 : 
                dists[i][j] = idist[unit[0]][unit[1]]
                dists[j][i] = dists[i][j]
            else :
                dists[i][j] = idist[deserters[j-1][0]][deserters[j-1][1]]
                dists[j][i] = dists[i][j]

    for order in pm(range(1, len(deserters)+1), len(deserters)) :
        dist = 0
        c = 0
        for dest in order :
            dist += dists[c][dest]
            c = dest
        dist += dists[c][0]
        minDist = min(dist, minDist)
    return minDist

graph = [list(map(int, input().strip().split())) for _ in range(n)]
deserters = []
c = ()
for x in range(n) :
    for y in range(n) :
        if graph[x][y] == -1 : 
            graph[x][y] = 0
            c = (x, y)
        elif graph[x][y] == 0 : deserters.append((x, y))
print(minDist(c, graph, deserters))