[algorithm] 백준 알고리즘 1753번 문제 정답코드

# CONSTANT
INFINITY = 'INF'
 
#################### high-level module ####################
 
def dijkstra(g, src):
    dist = [None] # distance list
    pq = [] # vertex pq (distance priority)
 
    for vtx in g.vtx_list[1:]:
        dist.append(INFINITY)
        pq.append(vtx)
    dist[src] = 0
 
    ### round pq ###
    while len(pq) != 0:
        vtx1 = remove_min_vtx(pq, dist)
        if dist[vtx1.idx] == INFINITY:
            break
        
        for edge in vtx1.adj_list: # 인접 정점들의 dist update
            vtx2 = edge.nxt
            new_dist = dist[vtx1.idx] + edge.w # src->vtx1 + u->vtx2
            if is_smaller(new_dist, dist[vtx2.idx]):
                dist[vtx2.idx] = new_dist
    
    return dist
 
#################### detail implementation ####################
 
# function of distance PQ
def remove_min_vtx(pq, dist):
    min_vtx = pq[0]
    for vtx in pq:
        if is_smaller(dist[vtx.idx], dist[min_vtx.idx]):
            min_vtx = vtx
    
    pq.remove(min_vtx)
    return min_vtx
 
# dist list function
def is_smaller(a, b):
    if a == INFINITY:
        return False
    elif b == INFINITY:
        return True
    else:
        return a < b
 
class DirGraph: # directed, weighted, adj_list graph
    def __init__(self, vtx_cnt):
        self.vtx_list = [None]
        
        for i in range(vtx_cnt):
            i=i+1
            self.add_vtx(i)
    
    def add_vtx(self, idx):
        self.vtx_list.append(Vertex(idx))
 
    def add_edge(self, u, v, w):
        edge = Edge(self.vtx_list[u], self.vtx_list[v], w)
        self.vtx_list[u].adj_list.append(edge) # adj_list == fwd_list
 
class Vertex:
    def __init__(self, idx):
        self.idx = idx
        self.adj_list = []
        
    def __str__(self):
        result = 'vtx_%d { ' % self.idx
        for edge in self.adj_list:
            result += '%s ' % edge
        return result+'}'
 
class Edge:
    def __init__(self, prev, nxt, weight):
        self.prev = prev
        self.nxt = nxt
        self.w = weight
        
    def __str__(self):
        return '(%d,%d,%d)' % (self.prev.idx, self.nxt.idx, self.w)
 
def main():
    vtx_cnt, edge_cnt = (int(i) for i in input().split())
    src = int(input())
    g = DirGraph(vtx_cnt)
    
    for i in range(edge_cnt):
        u, v, w = (int(i) for i in input().split())
        g.add_edge(u, v, w)
    
    dist = dijkstra(g, src)
    for i in dist:
        print(i)
 
def test_main():
    vtx_cnt, edge_cnt = 5,6
    src = 1
    g = DirGraph(vtx_cnt)
 
    g.add_edge(5,1,1)
    g.add_edge(1,2,2)
    g.add_edge(1,3,3)
    g.add_edge(2,3,4)
    g.add_edge(2,4,5)
    g.add_edge(3,4,6)
 
    print('<vtx_list>')
    for v in g.vtx_list[1:]:
        print(v)
    
    dist = dijkstra(g, src)
    print('\n<result>')
    for i in dist:
        print(i)
 
main()

주석 추가예정