diff --git a/dynamic_programming/floyd_warshall.py b/dynamic_programming/floyd_warshall.py index 2331f3e65483..9c874c2511d3 100644 --- a/dynamic_programming/floyd_warshall.py +++ b/dynamic_programming/floyd_warshall.py @@ -12,19 +12,56 @@ def __init__(self, n=0): # a graph with Node 0,1,...,N-1 ] # dp[i][j] stores minimum distance from i to j def add_edge(self, u, v, w): + """ + Adds a directed edge from node u to node v with weight w. + + >>> g = Graph(3) + >>> g.add_edge(0, 1, 5) + >>> g.dp[0][1] + 5 + """ self.dp[u][v] = w def floyd_warshall(self): + """ + Computes the shortest paths between all pairs of nodes using the Floyd-Warshall algorithm. + + >>> g = Graph(3) + >>> g.add_edge(0, 1, 1) + >>> g.add_edge(1, 2, 2) + >>> g.floyd_warshall() + >>> g.show_min(0, 2) + 3 + >>> g.show_min(2, 0) + inf + """ for k in range(self.n): for i in range(self.n): for j in range(self.n): self.dp[i][j] = min(self.dp[i][j], self.dp[i][k] + self.dp[k][j]) def show_min(self, u, v): + """ + Returns the minimum distance from node u to node v. + + >>> g = Graph(3) + >>> g.add_edge(0, 1, 3) + >>> g.add_edge(1, 2, 4) + >>> g.floyd_warshall() + >>> g.show_min(0, 2) + 7 + >>> g.show_min(1, 0) + inf + """ return self.dp[u][v] if __name__ == "__main__": + import doctest + + doctest.testmod() + + # Example usage graph = Graph(5) graph.add_edge(0, 2, 9) graph.add_edge(0, 4, 10) @@ -38,5 +75,5 @@ def show_min(self, u, v): graph.add_edge(4, 2, 4) graph.add_edge(4, 3, 9) graph.floyd_warshall() - graph.show_min(1, 4) - graph.show_min(0, 3) + print(graph.show_min(1, 4)) + print(graph.show_min(0, 3))