다익스트라를 사용하여 풀이하면 되는 문제입니다. 처음에는 순열을 사용하여 각 순서대로 다익스트라를 돌려 풀이하였는데, 당연하게도 순열 결과 배열의 길이 * 탈영병 인원수
만큼 다익스트라를 돌리다 보니 시간 초과가 떴습니다.
그래서 탈영병 인원만큼만 다익스트라를 돌려 각 탈영병들 간의 최소 비용을 미리 저장해둔 뒤, 순열을 사용하여 전체 최소 비용을 구할 때는 저장해둔 값을 불러와 계산하도록 변경하여 시간 초과를 해결하였습니다.
(탈영병 인원+1)*(탈영병 인원+1)
크기의 이중 배열을 생성한다. (0번째 인덱스는 부대라고 가정한다.)1, 2, ..., 탈영병_인원수
가 들어있는 배열의 가능한 순열을 계산한다.부대 -> 첫 번째 탈영병
+ 첫 번째 탈영병 -> 부대
의 비용부대 -> 첫 번째 탈영병
+ 첫 번째 탈영병 -> 두 번째 탈영병
+ ... + N-1 번째 탈영병 -> N 번째 탈영병
+ N 번째 탈영병 -> 부대
의 비용*순열, 최소 힙, 다익스트라를 모두 구현하여 코드가 조금 깁니다.
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))
*시간 초과로 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))