diff --git a/dijkstra.py b/dijkstra.py new file mode 100644 index 0000000000..03260a0f01 --- /dev/null +++ b/dijkstra.py @@ -0,0 +1,69 @@ +import heapq + +class Graph: + def __init__(self): + self.nodes = {} + + def add_node(self, value): + if value not in self.nodes: + self.nodes[value] = [] + + def add_edge(self, node1, node2, weight): + if node1 in self.nodes and node2 in self.nodes: + self.nodes[node1].append((node2, weight)) + else: + raise ValueError("One or both of the nodes are not present in the graph.") + + def dijkstra(self, source): + distances = {node: float('infinity') for node in self.nodes} + distances[source] = 0 + priority_queue = [(0, source)] + + while priority_queue: + current_distance, current_node = heapq.heappop(priority_queue) + + if current_distance > distances[current_node]: + continue + + for neighbor, weight in self.nodes[current_node]: + distance = current_distance + weight + + if distance < distances[neighbor]: + distances[neighbor] = distance + heapq.heappush(priority_queue, (distance, neighbor)) + + return distances + +def reconstruct_path(nodes_distances, end_node): + path = [end_node] + current_node = end_node + current_distance = nodes_distances[end_node] + + while current_distance != 0: + for neighbor, weight in graph.nodes[current_node]: + if nodes_distances[neighbor] == current_distance - weight: + path.append(neighbor) + current_node = neighbor + current_distance -= weight + break + + return path[::-1] + +if __name__ == "__main__": + graph = Graph() + graph.add_node("A") + graph.add_node("B") + graph.add_node("C") + graph.add_node("D") + + graph.add_edge("A", "B", 10) + graph.add_edge("A", "C", 15) + graph.add_edge("B", "D", 20) + graph.add_edge("C", "D", 25) + graph.add_edge("C", "B", 5) + + distances = graph.dijkstra("A") + print("Distances:", distances) + + path = reconstruct_path(distances, "D") + print("Shortest Path:", path) \ No newline at end of file diff --git a/test_dijkstra.py b/test_dijkstra.py new file mode 100644 index 0000000000..153d31017a --- /dev/null +++ b/test_dijkstra.py @@ -0,0 +1,79 @@ + import heapq +from typing import Dict, List, Optional + +class Graph: + """Represents a weighted graph with nodes and edges.""" + + def __init__(self): + self.nodes: Dict[str, List[tuple[str, int]]] = {} + + def add_node(self, node: str) -> None: + """Adds a node to the graph.""" + if node not in self.nodes: + self.nodes[node] = [] + + def add_edge( + self, node1: str, node2: str, weight: int, allow_duplicates: bool = False + ) -> None: + """Adds a weighted edge between two nodes.""" + if (node1 not in self.nodes) or (node2 not in self.nodes): + raise KeyError(f"Either {node1} or {node2} not found in the graph.") + + edge = (node2, weight) + + if not allow_duplicates and edge in self.nodes[node1]: + return + + self.nodes[node1].append(edge) + + def dijkstra(self, source: str) -> Dict[str, Optional[int]]: + """ + Implements Dijkstra's algorithm to find the shortest path from the source node to all other nodes. + + Returns a dictionary with node names as keys and shortest path distances as values. + """ + distances = {node: float("infinity") for node in self.nodes} + distances[source] = 0 + priority_queue = [(0, source)] + + while priority_queue: + current_distance, current_node = heapq.heappop(priority_queue) + + if current_distance > distances[current_node]: + continue + + for neighbor, edge_weight in self.nodes[current_node]: + distance = current_distance + edge_weight + + if distance < distances[neighbor]: + distances[neighbor] = distance + heapq.heappush(priority_queue, (distance, neighbor)) + + return distances + + +def test_dijkstra() -> None: + graph = Graph() + + # Add nodes and edges + graph.add_node("A") + graph.add_node("B") + graph.add_node("C") + graph.add_node("D") + + graph.add_edge("A", "B", 10) + graph.add_edge("A", "C", 5) + graph.add_edge("B", "D", 1) + graph.add_edge("C", "D", 20) + + # Test Dijkstra's algorithm + distances = graph.dijkstra("A") + + assert distances["A"] == 0 + assert distances["B"] == 10 + assert distances["C"] == 5 + assert distances["D"] == 6 + + +if __name__ == "__main__": + test_dijkstra() \ No newline at end of file