Facebook Hacker Cup Online Round 1a - After the Dance Battle

dance.py

import sys

# Facebook Hacker Cup Online Round 1a - After the Dance Battle
# http://www.facebook.com/hackercup/problems.php?round=144428782277390

def dijkstra(graph, start, end):    
    dist = {}
    previous = {}

    for v in graph.keys():
        dist[v] = float("inf")
        previous[v] = ""
        
    dist[start] = 0
    
    Q = graph.keys()
    
    while len(Q) > 0:
        u = Q[0]
        shortest = dist[Q[0]]
        
        for x in Q:
            if dist[x] < shortest:
                u = x
                shortest = dist[x]
        
        if dist[u] == float("inf"):
            break
        
        Q.remove(u)
        
        for v, child_value in graph[u].items():
            alt = dist[u] + child_value
            if alt < dist[v]:
                dist[v] = alt
                previous[v] = u
                
    path = []
    node = end
    
    while not (node == start):
        path.insert(0, node)
        node = previous[node]
    
    path.insert(0, start)
    return len(path)-1

def main():
    with open(sys.argv[1]) as fp:
        input = fp.read().split()
    
        index = 0
        n = int(input[index])
        index += 1
    
        for x in range(1,n+1):
            r = int(input[index])
            c = int(input[index+1])
            index += 2
            table = [[0]*c for z in range(r)]
      
            for i in range(r):
                for j in range(c):
                    y = input[index]
                    table[i][j] = y[j]
                index += 1
            
            graph = {}
            colors = {}
            for i in range(r):
                for j in range(c):
                    edges = {}
                    
                    # check colors
                    color = table[i][j]
                    if color >= '1' and color <= '9':
                        if not colors.has_key(color):
                            colors[color] = []
                        colors[color].append((i,j))
                    
                    # check left
                    if j > 0:
                        if table[i][j-1] != 'W':
                            edges[(i,j-1)] = 1
                    
                    # check right
                    if j < c-1:
                        if table[i][j+1] != 'W':
                            edges[(i,j+1)] = 1
                    
                    # check down
                    if i < r-1:
                        if table[i+1][j] != 'W':
                            edges[(i+1,j)] = 1
                    
                    # check up
                    if i > 0:
                        if table[i-1][j] != 'W':
                            edges[(i-1,j)] = 1
                    
                    graph[(i,j)] = edges
            
            for a, b in graph.keys():
                color = table[a][b]
                if color >= '1' and color <= '9':
                    edges = graph[(a,b)]
                    for y, z in colors[color]:
                        edges[(y,z)] = 1
            
            for j in range(c):
                if table[0][j] == 'S':
                    start = (0,j)
                
            for j in range(c):
                if table[r-1][j] == 'E':
                    end = (r-1, j)
            
            distance = dijkstra(graph, start, end)
            print distance

if __name__ == "__main__":
    main()