백준 최단경로 (Python/Go)

문제 링크: www.acmicpc.net/problem/1753

1753번: 최단경로

 

기본 다익스트라 문제입니다.

아래는 파이썬 코드입니다.

from collections import defaultdict
import heapq
 
 
def dijkstra(start):
    dist = defaultdict(int)
    heap = [(0, start)]
    while heap:
        weight, node = heapq.heappop(heap)
        if node not in dist:
            dist[node] = weight
            for next_node, next_weight in graph[node]:
                cost = weight + next_weight
                heapq.heappush(heap, (cost, next_node))
 
    return dist
 
 
n, e = map(int, input().split())
k = int(input())
graph = defaultdict(list)
for _ in range(e):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))
 
dist = dijkstra(k)
 
for i in range(1, n+1):
    if i in dist:
        print(dist[i])
    else:
        print('INF')

 

아래는 Go코드입니다.

package main
 
import (
    "bufio"
    "container/heap"
    "fmt"
    "os"
    "strconv"
    "strings"
)
 
type Item struct {  // graph[u] = [(w, v) ...] 형태로 저장하기 위해 + heap에서도 쓰임
    node   int
    weight int
}
 
type PriorityQueue []*Item
 
func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool {
    return pq[i].weight < pq[j].weight
}
 
func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
}
 
func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    old[n-1] = nil  // avoid memory leak
    *pq = old[0 : n-1]
    return item
}
 
func (pq *PriorityQueue) Push(x interface{}) {
    item := x.(*Item)
    *pq = append(*pq, item)
}
 
 
func GetValues(input string) (int, int, int) {  // u, v, w 값을 int로 받기 위한 함수
    var nums []int
    for _, numString := range strings.Fields(input) {
        num, _ := strconv.Atoi(numString)
        nums = append(nums, num)
    }
    return nums[0], nums[1], nums[2]
}
 
func dijkstra(startNode int, graph map[int][]Item) map[int]int {
    dist := make(map[int]int)
    pq := PriorityQueue{&Item{startNode, 0}}
    for {
        if len(pq) == 0 {
            break
        }
 
        curInfo := heap.Pop(&pq).(*Item)
        curNode, curWeight := curInfo.node, curInfo.weight
        if _, ok := dist[curNode]; !ok {  // !ok -> curNode가 dist에 없으
            dist[curNode] = curWeight
            for _, nextInfo := range graph[curNode] {
                nextNode, nextWeight := nextInfo.node, nextInfo.weight
                totalWeight := curWeight + nextWeight
                heap.Push(&pq, &Item{nextNode, totalWeight})
            }
        }
    }
    return dist
}
 
func main() {
    var n, e int
    fmt.Scanln(&n, &e)
 
    var k int
    fmt.Scanln(&k)
 
    scanner := bufio.NewScanner(os.Stdin)
    graph := make(map[int][]Item)
    var u, v, w int
 
    for i:=0; i<e; i++ {
        scanner.Scan()
        input := scanner.Text()
        u, v, w = GetValues(input)
        graph[u] = append(graph[u], Item{v, w})
    }
 
    dist := dijkstra(k, graph)
 
    for i:=1; i<n+1; i++ {
        if _, ok := dist[i]; ok {  // ok에 걸리면 dist에 거리가 저장되어 있다는 뜻
            fmt.Println(dist[i])
        } else {
            fmt.Println("INF")
        }
    }
}

파이썬의 heapq 모듈에 비해 Go에서는 heap을 적용하는 과정이 복잡했습니다.

속도는 Go는 504ms, 파이썬은 1252ms가 나왔습니다. 그러나 코드길이는 파이썬이 3배가량 짧습니다.