Monday, February 17, 2020

USACO Bessie Come Home, Dijkstra's Algorithm (Python)

Bessie Come Home

USACO Bessie Come Home

In [1]:
# SAMPLE INPUT
# 5
# A d 6
# B d 3
# C e 9
# d Z 8
# e Z 3
In [2]:
fin = open ('comehome.in', 'r')
fout = open ('comehome.out', 'w')
P = int(fin.readline().strip())
print(P)
G = dict()
for i in range(P):
    a, b, d = fin.readline().strip().split()
    print(a,b,d)
    if a not in G:
        G[a]=[]
    if b not in G:
        G[b]=[]
                        
    G[a].append((b, int(d)))
    G[b].append((a, int(d)))
5
A d 6
B d 3
C e 9
d Z 8
e Z 3
In [3]:
for e in G.items():
    print(e)
('A', [('d', 6)])
('d', [('A', 6), ('B', 3), ('Z', 8)])
('B', [('d', 3)])
('C', [('e', 9)])
('e', [('C', 9), ('Z', 3)])
('Z', [('d', 8), ('e', 3)])
In [4]:
capitals = "ABCDEFGHIJKLMNOPQRSTUVWXY"
BIGCOST = 1e12

dist = {v: BIGCOST for v in G}
dist['Z'] = 0
visited = set()

while len(visited) < len(G):
    unvisited = [(v,d) for v,d in dist.items() if v not in visited]
    a,d = min(unvisited,key=lambda x: x[1])
    if a in capitals:
        break
    visited.add(a)
    
    for b,d in G[a]:
        if dist[b] > dist[a]+d:
            dist[b] = dist[a]+d
     
print(a,d)
B 11
In [5]:
fout.write(a+" "+str(d)+"\n")
fout.close()