| 2 | |
| 3 | |
| 4 | class Graph: |
| 5 | def __init__(self, n=0): # a graph with Node 0,1,...,N-1 |
| 6 | self.n = n |
| 7 | self.w = [ |
| 8 | [math.inf for j in range(n)] for i in range(n) |
| 9 | ] # adjacency matrix for weight |
| 10 | self.dp = [ |
| 11 | [math.inf for j in range(n)] for i in range(n) |
| 12 | ] # dp[i][j] stores minimum distance from i to j |
| 13 | |
| 14 | def add_edge(self, u, v, w): |
| 15 | """ |
| 16 | Adds a directed edge from node u |
| 17 | to node v with weight w. |
| 18 | |
| 19 | >>> g = Graph(3) |
| 20 | >>> g.add_edge(0, 1, 5) |
| 21 | >>> g.dp[0][1] |
| 22 | 5 |
| 23 | """ |
| 24 | self.dp[u][v] = w |
| 25 | |
| 26 | def floyd_warshall(self): |
| 27 | """ |
| 28 | Computes the shortest paths between all pairs of |
| 29 | nodes using the Floyd-Warshall algorithm. |
| 30 | |
| 31 | >>> g = Graph(3) |
| 32 | >>> g.add_edge(0, 1, 1) |
| 33 | >>> g.add_edge(1, 2, 2) |
| 34 | >>> g.floyd_warshall() |
| 35 | >>> g.show_min(0, 2) |
| 36 | 3 |
| 37 | >>> g.show_min(2, 0) |
| 38 | inf |
| 39 | """ |
| 40 | for k in range(self.n): |
| 41 | for i in range(self.n): |
| 42 | for j in range(self.n): |
| 43 | self.dp[i][j] = min(self.dp[i][j], self.dp[i][k] + self.dp[k][j]) |
| 44 | |
| 45 | def show_min(self, u, v): |
| 46 | """ |
| 47 | Returns the minimum distance from node u to node v. |
| 48 | |
| 49 | >>> g = Graph(3) |
| 50 | >>> g.add_edge(0, 1, 3) |
| 51 | >>> g.add_edge(1, 2, 4) |
| 52 | >>> g.floyd_warshall() |
| 53 | >>> g.show_min(0, 2) |
| 54 | 7 |
| 55 | >>> g.show_min(1, 0) |
| 56 | inf |
| 57 | """ |
| 58 | return self.dp[u][v] |
| 59 | |
| 60 | |
| 61 | if __name__ == "__main__": |