Ruigu Electronic

from collections import deque
import heapq

class WeightedGraph:
    def __init__(self, directed=False):
        self.graph = {}  # Adjacency list: {vertex: [(neighbor, weight), ...]}
        self.directed = directed  # True = weighted digraph, False = undirected weighted graph

    # -------------------------- Core Operations --------------------------
    def add_vertex(self, vertex):
        """Add a vertex (no duplicates)"""
        if vertex not in self.graph:
            self.graph[vertex] = []
            return True
        return False

    def add_edge(self, u, v, weight):
        """Add a weighted edge between u and v (directional if directed=True)"""
        if u not in self.graph or v not in self.graph:
            raise ValueError("Both vertices must exist in the graph")
        # Validate weight (allow integers/floats)
        if not isinstance(weight, (int, float)):
            raise TypeError("Weight must be a number")
        # Add edge u→v
        self.graph[u].append((v, weight))
        # Add reverse edge v→u if undirected
        if not self.directed:
            self.graph[v].append((u, weight))

    # -------------------------- Shortest Path Algorithms --------------------------
    def dijkstra(self, start_vertex):
        """
        Shortest path from start_vertex to all others (non-negative weights only)
        Returns: {vertex: shortest_distance, ...}
        """
        if start_vertex not in self.graph:
            raise ValueError("Start vertex not in graph")
        
        # Initialize distances: infinity except start_vertex (0)
        INF = float('inf')
        distance = {vertex: INF for vertex in self.graph}
        distance[start_vertex] = 0
        
        # Priority queue: (current_distance, current_vertex) → always expands shortest known path
        priority_queue = [(0, start_vertex)]
        visited = set()

        while priority_queue:
            current_dist, u = heapq.heappop(priority_queue)
            
            # Skip if vertex already processed (we found a shorter path earlier)
            if u in visited:
                continue
            visited.add(u)

            # Relax all outgoing edges from u
            for (v, weight) in self.graph[u]:
                if weight < 0:
                    raise ValueError("Dijkstra's Algorithm does not support negative weights")
                if distance[v] > distance[u] + weight:
                    distance[v] = distance[u] + weight
                    heapq.heappush(priority_queue, (distance[v], v))

        return distance

    def bellman_ford(self, start_vertex):
        """
        Shortest path from start_vertex to all others (supports negative weights, no negative cycles)
        Returns: {vertex: shortest_distance, ...} or raises error if negative cycle exists
        """
        if start_vertex not in self.graph:
            raise ValueError("Start vertex not in graph")
        
        V = len(self.graph)
        INF = float('inf')
        distance = {vertex: INF for vertex in self.graph}
        distance[start_vertex] = 0

        # Step 1: Relax all edges V-1 times (longest shortest path has V-1 edges)
        edges = self.get_edges()
        for _ in range(V - 1):
            updated = False
            for (u, v, weight) in edges:
                if distance[u] != INF and distance[v] > distance[u] + weight:
                    distance[v] = distance[u] + weight
                    updated = True
            if not updated:
                break  # No more updates → early exit

        # Step 2: Check for negative weight cycles
        for (u, v, weight) in edges:
            if distance[u] != INF and distance[v] > distance[u] + weight:
                raise ValueError("Graph contains a negative weight cycle (shortest path undefined)")

        return distance

    def floyd_warshall(self):
        """
        All-pairs shortest path (supports negative weights, no negative cycles)
        Returns: Distance matrix {u: {v: shortest_distance, ...}, ...}
        """
        vertices = list(self.graph.keys())
        V = len(vertices)
        INF = float('inf')

        # Initialize distance matrix: dist[i][j] = weight of edge i→j, else INF
        dist = {u: {v: INF for v in vertices} for u in vertices}
        for u in vertices:
            dist[u][u] = 0  # Distance from vertex to itself is 0
            for (v, weight) in self.graph[u]:
                dist[u][v] = weight

        # Step 1: Update distances using intermediate vertices
        for k in vertices:  # k = intermediate vertex
            for i in vertices:  # i = source
                for j in vertices:  # j = target
                    if dist[i][k] != INF and dist[k][j] != INF:
                        if dist[i][j] > dist[i][k] + dist[k][j]:
                            dist[i][j] = dist[i][k] + dist[k][j]

        # Step 2: Check for negative weight cycles (distance from vertex to itself < 0)
        for u in vertices:
            if dist[u][u] < 0:
                raise ValueError("Graph contains a negative weight cycle")

        return dist

    # -------------------------- Minimum Spanning Tree (MST) Algorithms --------------------------
    def prim_mst(self):
        """
        Prim's Algorithm: MST for undirected, connected weighted graphs
        Returns: List of MST edges [(u, v, weight), ...] and total MST weight
        """
        if self.directed:
            raise ValueError("Prim's Algorithm is for undirected graphs")
        if not self.is_connected():
            raise ValueError("Graph must be connected to have an MST")
        
        vertices = list(self.graph.keys())
        if not vertices:
            return [], 0
        
        # Initialize: start with first vertex
        visited = set([vertices[0]])
        mst_edges = []
        total_weight = 0

        # Priority queue: (edge_weight, u, v) → always pick the smallest edge connecting MST to non-MST
        priority_queue = []
        for (v, weight) in self.graph[vertices[0]]:
            heapq.heappush(priority_queue, (weight, vertices[0], v))

        while priority_queue and len(visited) < len(vertices):
            weight, u, v = heapq.heappop(priority_queue)
            if v not in visited:
                visited.add(v)
                mst_edges.append((u, v, weight))
                total_weight += weight

                # Add edges from v to non-MST vertices
                for (neighbor, w) in self.graph[v]:
                    if neighbor not in visited:
                        heapq.heappush(priority_queue, (w, v, neighbor))

        return mst_edges, total_weight

    # -------------------------- Helper Methods --------------------------
    def get_edges(self):
        """Return all edges as (u, v, weight) tuples"""
        edges = []
        for u in self.graph:
            for (v, weight) in self.graph[u]:
                # Avoid duplicate edges in undirected graphs (store (u, v) only if u < v)
                if not self.directed and u > v:
                    continue
                edges.append((u, v, weight))
        return edges

    def is_connected(self):
        """Check if undirected graph is connected (BFS-based)"""
        if self.directed:
            raise ValueError("is_connected() is for undirected graphs")
        if not self.graph:
            return True
        
        start_vertex = next(iter(self.graph.keys()))
        visited = set()
        queue = deque([start_vertex])
        visited.add(start_vertex)

        while queue:
            u = queue.popleft()
            for (v, _) in self.graph[u]:
                if v not in visited:
                    visited.add(v)
                    queue.append(v)

        return len(visited) == len(self.graph)

    def __str__(self):
        """String representation (adjacency list with weights)"""
        result = []
        for u in self.graph:
            neighbors = ", ".join([f"{v} (w={weight})" for (v, weight) in self.graph[u]])
            direction = "→" if self.directed else "↔"
            result.append(f"{u} {direction} [{neighbors}]")
        return "\n".join(result)

if __name__ == "__main__":
    # Test 1: Undirected Weighted Graph (Road Network)
    road_network = WeightedGraph(directed=False)
    cities = ["NYC", "LA", "Chicago", "Houston", "Miami"]
    for city in cities:
        road_network.add_vertex(city)
    # Add edges (city1, city2, distance in miles)
    road_network.add_edge("NYC", "Chicago", 800)
    road_network.add_edge("NYC", "Miami", 1300)
    road_network.add_edge("Chicago", "LA", 2000)
    road_network.add_edge("Chicago", "Houston", 1100)
    road_network.add_edge("Houston", "LA", 1500)
    road_network.add_edge("Houston", "Miami", 1200)

    print("=== Undirected Weighted Graph (Road Network) ===")
    print(road_network)

    # Dijkstra's Algorithm (shortest path from NYC)
    print("\n=== Dijkstra's Shortest Path from NYC ===")
    shortest_paths = road_network.dijkstra("NYC")
    for city, dist in shortest_paths.items():
        print(f"NYC → {city}: {dist} miles")  # Output: NYC→LA: 2800 (NYC→Chicago→LA)

    # Prim's MST (minimum cost to connect all cities)
    print("\n=== Prim's MST (Road Network) ===")
    mst_edges, total_weight = road_network.prim_mst()
    print("MST Edges:", mst_edges)
    print("Total MST Weight:", total_weight)  # Output: 800+1100+1200+1500 = 4600

    # Test 2: Directed Weighted Graph (Network Latency)
    network = WeightedGraph(directed=True)
    nodes = ["A", "B", "C", "D"]
    for node in nodes:
        network.add_vertex(node)
    # Add edges (node1→node2, latency in ms)
    network.add_edge("A", "B", 5)
    network.add_edge("A", "C", 10)
    network.add_edge("B", "C", 2)
    network.add_edge("B", "D", 15)
    network.add_edge("C", "D", 3)
    network.add_edge("D", "C", -1)  # Negative weight (no cycle → Bellman-Ford works)

    print("\n=== Directed Weighted Graph (Network Latency) ===")
    print(network)

    # Bellman-Ford Algorithm (supports negative weights)
    print("\n=== Bellman-Ford Shortest Path from A ===")
    try:
        bf_paths = network.bellman_ford("A")
        for node, latency in bf_paths.items():
            print(f"A→{node}: {latency} ms")  # Output: A→D: 5+2+3=10 ms
    except ValueError as e:
        print(e)

    # Floyd-Warshall (All-Pairs Shortest Path)
    print("\n=== Floyd-Warshall All-Pairs Shortest Path ===")
    try:
        fw_matrix = network.floyd_warshall()
        for u in fw_matrix:
            for v in fw_matrix[u]:
                print(f"{u}→{v}: {fw_matrix[u][v]} ms", end=" | ")
            print()
    except ValueError as e:
        print(e)