Weighted_Graph.py

import heapq

class Graph:
    
    def __init__(self, directed = False):
        self.graph = {}
        self.directed = directed
        
        
    def add_node(self, node):
        if node not in self.graph:
            self.graph[node] = []
            
    def add_edge(self, node1, node2, weight):
        
        if node1 in self.graph and node2 in self.graph:
            self.graph[node1].append((node2, weight))
            
            if not self.directed:
                self.graph[node2].append((node1, weight))
                
    def remove_edge(self, node1, node2):
        
        if node1 in self.graph and node2 in self.graph:
            self.graph[node1] = [(n, w) for n, w in self.graph[node1] if n!= node2]
            
            if not self.directed:
                self.graph[node2] = [(n, w) for n, w in self.graph[node2] if n != node1]
                
    def remove_node(self, node):
        
        if node in self.graph:
            
            for neighbor, _ in self.graph[node]:
                self.graph[neighbor] = [(n, w) for n, w in self.graph[neighbor] if n != node]
                
            del self.graph[node]
            
    def BFS(self, start):
        
        visited = set()
        queue = [start]
        result = []
        
        while queue:
            
            node = queue.pop(0)
            if node not in visited:
                visited.add(node)
                result.append(node)
                
                for neighbor, _ in self.graph[node]:
                    if neighbor not in visited:
                        queue.append(neighbor)
                        
        return result
    
    def DFS(self, start):
        
        visited = set()
        stack = [start]
        result = []
        
        while stack:
            node = stack.pop()
            
            if node not in visited:
                visited.add(node)
                result.append(node)
                
                for neighbor, _ in reversed(self.graph[node]):
                    if neighbor not in visited:
                        stack.append(neighbor)
                        
        return result
    
    def djkstra(self, start):
        
        if start not in self.graph:
            return {}
        
        distances = {node: float('inf') for node in self.graph}
        distances[start] = 0
        heap = [(0, start)]
        
        while heap:
            
            curr_dist, curr_node = heapq.heappop(heap)
            
            if curr_dist > distances[curr_node]:
                continue
            
            for neighbor, weight in self.graph[curr_node]:
                
                distance = curr_dist + weight
                
                if distance < distances[neighbor]:
                    distances[neighbor] = distance
                    heapq.heappush(heap, (distance, neighbor))
                    
        return distances
            
    def display(self):
        for node, edges in self.graph.items():
            print(f"{node}: {edges}")
            
            
            
if __name__ == "__main__":
    g = Graph()
    g.add_node("A")
    g.add_node("B")
    g.add_node('C')
    g.add_node('D')
    
    
    g.add_edge('A', 'B', 5)
    g.add_edge('A', 'C', 3)
    g.add_edge('B', 'C', 2)
    g.add_edge('A', 'D', 4)
    
    g.remove_edge('A', 'C')
    
    g.display()
    
    print(g.BFS('A'))
    print(g.DFS('A'))
    print(g.djkstra('A'))