hacker-cup/2017.1-RoundOne/c.ManiacMoving.py

import numpy as np

def directDist(i, j, ABG):
    dist = -1
    for A, B, G in ABG:
        if ((A == i and B == j) or (A == j and B == i)) and (dist == -1 or G < dist):
            dist = G
    return dist

# From http://stackoverflow.com/a/22899400
def shortest(i, j, ABG, nodes):
    # print("i={0}, j={1}, ABG={2}, nodes={3}".format(i, j, ABG, nodes))
    distances = {}
    for A, B, G in ABG:
        # print(distances)
        if A not in distances or B not in distances:
            distances[A] = {B: G}
            distances[B] = {A: G}
            continue
        if B not in distances[A]:
            distances[A][B] = G
            distances[B][A] = G
            continue
        if G < distances[A][B]:
            distances[A][B] = G
            distances[B][A] = G

    unvisited = {node: None for node in nodes} #using None as +inf
    visited = {}
    current = i
    currentDistance = 0
    unvisited[current] = currentDistance

    while True:
        for neighbour, distance in distances[current].items():
            if neighbour not in unvisited: continue
            newDistance = currentDistance + distance
            if unvisited[neighbour] is None or unvisited[neighbour] > newDistance:
                unvisited[neighbour] = newDistance
        visited[current] = currentDistance
        if current == j:
            return visited
        del unvisited[current]
        if not unvisited: break
        candidates = [node for node in unvisited.items() if node[1]]
        # print("unvisited", unvisited)
        # print("candidates", candidates)
        current, currentDistance = sorted(candidates, key = lambda x: x[1])[0]

    return visited

def routes(ABG, SD):
    r = {}
    townsInSD = set()
    townsInABG = set()
    # We only compute the routes for towns in SD
    for S, D in SD:
        # print("S={0}, D={1}".format(S, D))
        townsInSD.add(S)
        townsInSD.add(D)
    for A, B, G in ABG:
        # print("A={0}, B={1}, G={2}".format(A, B, G))
        townsInABG.add(A)
        townsInABG.add(B)
    # print(townsInSD)
    # print(townsInABG)
    if len(townsInABG) < len(townsInSD):
        return {}
    allTowns = townsInSD | townsInABG

    # print("Shortest:", shortest(1, 2, ABG, allTowns))
    for i in townsInSD:
        for j in townsInSD:
            if (i,j) in r:
                continue
            shortst = shortest(i, j, ABG, allTowns)
            print("Shortest:", shortst)
            dist = shortst[j]
            r[(i,j)] = dist
            r[(j,i)] = dist
    return r

def solve(N, M, K, ABG, SD):
    print("N={0}, M={1}, K={2}, ABG={3}, SD={4}".format(N, M, K, ABG, SD))
    if K <= 0:
        return 0
    r = routes(ABG, SD)
    print(r)
    if r == {}:
        return -1
    return 'gnagna'

if __name__ == "__main__":
    import fileinput
    f = fileinput.input()

    T = int(f.readline())
    for case in range(1, T+1):
        N, M, K = [int(i) for i in f.readline().split()]
        ABG = [[int(i) for i in f.readline().split()] for x in range(M)]
        SD = [[int(i) for i in f.readline().split()] for x in range(K)]
        solution = solve(N, M, K, ABG, SD)
        print("Case #{0}: {1}".format(case, solution))