Programming Challenge 54 - Fastest Railroad Route [Python]

import sys, time, copy, heapq, pprint

STATIONS = set()
COSTS = {'+':1, 'S':3}
INF = 1000000000

def dijkstra(G, start, end):
    def flatten(L):
        while len(L) > 0:
            yield L[0]
            L = L[1]
    q = [(0, start, ())]
    visited = set()
    while True:
        (cost, v1, path) = heapq.heappop(q)
        if v1 not in visited:
            visited.add(v1)
            if v1 == end:
                return list(flatten(path))[::-1] + [v1]
            path = (v1, path)
            for (v2, cost2) in G[v1].iteritems():
                if v2 not in visited:
                    heapq.heappush(q, (cost + cost2, v2, path))

def hamiltonian(dist):
    n = len(dist)
    dp = [[INF for _ in range(n)] for _ in range(1 << n)]
    for i in range(n): dp[1 << i][i] = 0
    for mask in range(1<<n):
        for i in range(n):
            if mask & (1 << i):
                for j in range(n):
                    if mask & (1 << j):
                        dp[mask][i] = min(dp[mask][i], dp[mask ^ 1 << i][j] + dist[j][i])
    res = INF
    v   = 0
    for i in range(n):
        value = dp[(1 << n) - 1][i]
        if value < res:
            res = value
            v = i

    order = n*[0]
    cur = (1 << n) - 1
    cur ^= 1 << v
    rest = res
    order[n - 1] = v
    for i in range(n - 2, -1, -1):
        for j in range(n):
            if cur & 1 << j:
               if rest == dp[cur][j] + dist[j][v]:
                   rest -= dist[j][v]
                   v = j
                   order[i] = v
                   cur ^= 1 << v
                   break
    return res, order

def print_track(content):
    s1 = s2 = ' '*3
    last = '0'
    for i in range(len(content[0])):
        s = str(i).zfill(2)
        if s[0] != last:
            s1 += s[0]
            last = s[0]
        else:
            s1 += ' '
        s2 += s[1]
    print s1
    print s2
    no = 0
    for line in content:
        print str(no).zfill(2), line
        no += 1

def get_nexts(data, x, y, visited, stack, STATIONS):
    res = []
    moves = [(x,y-1), (x+1,y-1), (x+1,y), (x+1,y+1), (x,y+1), (x-1,y+1), (x-1,y), (x-1,y-1)]
    for cx,cy in moves:
        if (cx,cy) not in visited:
            d = data[cy][cx]
            if d in STATIONS or d == 'S':
                res.append((data[cy][cx],cx,cy))
                break
    if not len(res):
        for cx,cy in moves:
            if (cx,cy) not in visited and (cx,cy) not in stack:
               d = data[cy][cx]
               if d != '.':
                   res.append((data[cy][cx],cx,cy))
    return res

def bfs(data, station, x, y, STATIONS):
    g = {}
    stack = set()
    q = [((x,y), station, [])]
    stack.add((x,y))
    visited = set()
    while len(q):
        (x, y), vertex, l = q.pop(0)
        if (x, y) not in visited:
            visited.add((x, y))
            for (v2, x2, y2) in get_nexts(data, x, y, visited, stack, STATIONS):
                new_len = l + [v2]
                if v2 != '+':
                    if v2 == 'S':
                        new_vertex = 'S%d_%d'%(x2,y2)
                    else:
                        new_vertex = v2
                    g.setdefault(vertex, {})[new_vertex] = sum([COSTS[e] for e in new_len])
                    vertex = new_vertex
                    new_len = []
                if (x2, y2) not in visited:
                    q.append(((x2, y2), vertex, new_len))
                    stack.add((x2,y2))
    return g

def find_coordinates(data, station):
    x = y = 0
    for row in data:
        if station in row:
           x = row.index(station)
           break
        else:
           y += 1
    return x, y

def main():
    start = time.time()
    content = open(sys.argv[1]).read().split()

    print_track(content)

    # find all stations in text file
    STATIONS = set()
    for row in content:
        for e in row:
            if e not in ['+','.','S']:
               STATIONS.add(e)

    for s in STATIONS: COSTS[s]=5
    data = [list(content[i]) for i in range(len(content))]

    # build the graph using bfs
    G = {}
    for station in sorted(STATIONS):
        x, y = find_coordinates(data, station)
        SG = bfs(data, station, x, y, STATIONS)
        for k,v in SG.items():
            for sv in v:
                 value = SG[k][sv]
                 G.setdefault(k,{})[sv] = value

    N = len(G.keys())
    d = [[INF for _ in xrange(N)] for _ in xrange(N)]

    SYMBOLS = G.keys()
    for k,v in G.items():
        for sv in v:
            d[SYMBOLS.index(k)][SYMBOLS.index(sv)] = G[k][sv]

    # Floyd Warshall
    floyd = copy.deepcopy(d)
    for k in xrange(N):
        for i in xrange(N):
            for j in xrange(N):
                if i != j:
                    floyd[i][j] = min(floyd[i][j], floyd[i][k] + floyd[k][j])

    # calc paths between only all stations
    N = len(STATIONS)
    STATION_LIST = list(STATIONS)
    m = [[INF for _ in xrange(N)] for _ in xrange(N)]
    for src in STATIONS:
        for dst in STATIONS:
            if src != dst:
                m[STATION_LIST.index(src)][STATION_LIST.index(dst)] = floyd[SYMBOLS.index(src)][SYMBOLS.index(dst)]

    # calc shortest hamiltonian path
    shortest_path_len, path = hamiltonian(m)
    real_path = []

    for i in range(0,len(path)-1):
        src = STATION_LIST[path[i]]
        dst = STATION_LIST[path[i+1]]
        p = dijkstra(G, src, dst)
        calc_size = 0
        for i in range(len(p)-1):
            calc_size += G[p[i]][p[i+1]]
        real_path.append((src, dst, calc_size, p))

    for src, dst, calc_size, p in real_path:
        print src,' -> ',dst,' = ', calc_size, ' path = ', p

    print "Shortest path = ", shortest_path_len
    print "Time (s):", time.time()-start

main()