tester.py

#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "with python36Packages; [networkx numpy python3]"

import networkx as nx
import numpy as np
import sys

def main():
    assert(len(sys.argv) == 3)

    with open(sys.argv[1]) as f:
        it = iter(f)
        next(it)
        l = []
        for line in f:
            l.append(list(map(int, line.split())))

        xs = np.array(l)

    G = nx.Graph()

    with open(sys.argv[2]) as f:
        for (i, line) in enumerate(f):
            u, v, d = map(int, line.split())
            G.add_edge(u - 1, v - 1, key=i, weight=d)

    ys = nx.floyd_warshall_numpy(G, range(i + 1)).astype(int)

    if (xs == ys).all():
        print("Correct!")
    else:
        print("Wrong!")


if __name__ == "__main__":
    main()