diff --git a/graphs/graphs_floyd_warshall.py b/graphs/graphs_floyd_warshall.py index aaed9ac5df8b..c14053e7a2c9 100644 --- a/graphs/graphs_floyd_warshall.py +++ b/graphs/graphs_floyd_warshall.py @@ -6,13 +6,14 @@ def _print_dist(dist, v): - print("\nThe shortest path matrix using Floyd Warshall algorithm\n") + print("The shortest path matrix using Floyd Warshall algorithm\n") for i in range(v): for j in range(v): + end_char = "" if j == v - 1 else " " if dist[i][j] != float("inf"): - print(int(dist[i][j]), end="\t") + print(int(dist[i][j]), end=end_char) else: - print("INF", end="\t") + print("INF", end=end_char) print() @@ -31,6 +32,27 @@ def floyd_warshall(graph, v): 4. The above is repeated for each vertex k in the graph. 5. Whenever distance[i][j] is given a new minimum value, next vertex[i][j] is updated to the next vertex[i][k]. + + + >>> graph = [ + ... [0, 3, float('inf')], + ... [2, 0, float('inf')], + ... [float('inf'), 7, 0] + ... ] + + >>> expected = [ + ... [0, 3, float('inf')], + ... [2, 0, float('inf')], + ... [9, 7, 0] + ... ] + >>> dist, _ = floyd_warshall(graph, 3) + The shortest path matrix using Floyd Warshall algorithm + + 0 3 INF + 2 0 INF + 9 7 0 + >>> dist == expected + True """ dist = [[float("inf") for _ in range(v)] for _ in range(v)] diff --git a/graphs/tests/test_graphs_floyd_warshall.py b/graphs/tests/test_graphs_floyd_warshall.py new file mode 100644 index 000000000000..ef8bb230f1c9 --- /dev/null +++ b/graphs/tests/test_graphs_floyd_warshall.py @@ -0,0 +1,40 @@ +import pytest + +from graphs.graphs_floyd_warshall import floyd_warshall + + +def test_no_edges(): + graph = [ + [0, float("inf"), float("inf")], + [float("inf"), 0, float("inf")], + [float("inf"), float("inf"), 0], + ] + expected = [ + [0, float("inf"), float("inf")], + [float("inf"), 0, float("inf")], + [float("inf"), float("inf"), 0], + ] + dist, _ = floyd_warshall(graph, 3) + assert dist == expected + + +def test_with_edges(): + graph = [[0, 3, float("inf")], [2, 0, float("inf")], [float("inf"), 7, 0]] + expected = [[0, 3, float("inf")], [2, 0, float("inf")], [9, 7, 0]] + dist, _ = floyd_warshall(graph, 3) + assert dist == expected + + +def test_unreachable_vertices(): + graph = [ + [0, 1, float("inf")], + [float("inf"), 0, 2], + [float("inf"), float("inf"), 0], + ] + expected = [[0, 1, 3], [float("inf"), 0, 2], [float("inf"), float("inf"), 0]] + dist, _ = floyd_warshall(graph, 3) + assert dist == expected + + +if __name__ == "__main__": + pytest.main()