CS412-Applied-Algorithms/Railroad-Construction/cs412_railroad_a.py

"""
name: Nicholas Tamassia

Honor Code and Acknowledgments:

        This work complies with the JMU Honor Code.

       Comments here on your code and submission.
"""

from collections import defaultdict
from math import sqrt

type Coordinate = tuple[float, float]
type CoordMap = dict[int, Coordinate]
type WeightedGraph = dict[int, dict[int, float]]
type Edge = tuple[int, int, float]


def distance(a: Coordinate, b: Coordinate) -> float:
    x1, y1 = a
    x2, y2 = b
    return sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)


def create_graph(coords: CoordMap) -> WeightedGraph:
    graph: WeightedGraph = defaultdict(dict)

    coord_items = coords.items()

    for label1, coord1 in coord_items:
        for label2, coord2 in coord_items:
            if label1 == label2:
                continue

            weight = distance(coord1, coord2)
            graph[label1][label2] = weight
            graph[label2][label1] = weight

    return graph


def calculate_cost(graph: WeightedGraph) -> float:
    parent: dict[int, int] = {v: v for v in graph}
    rank: dict[int, int] = {v: 0 for v in graph}

    def find(v: int) -> int:
        if parent[v] != v:
            parent[v] = find(parent[v])
        return parent[v]

    def union(u: int, v: int) -> None:
        root_u, root_v = find(u), find(v)
        if root_u == root_v:
            return
        if rank[root_u] < rank[root_v]:
            parent[root_u] = root_v
        elif rank[root_u] > rank[root_v]:
            parent[root_v] = root_u
        else:
            parent[root_v] = root_u
            rank[root_u] += 1

    num_components = len(graph)
    mst_edges: list[Edge] = []

    while num_components > 1:
        cheapest: dict[int, Edge] = {}

        for u in graph:
            for v, w in graph[u].items():
                set_u = find(u)
                set_v = find(v)
                if set_u == set_v:
                    continue
                if set_u not in cheapest or cheapest[set_u][2] > w:
                    cheapest[set_u] = (u, v, w)
                if set_v not in cheapest or cheapest[set_v][2] > w:
                    cheapest[set_v] = (v, u, w)

        for edge in cheapest.values():
            u, v, w = edge
            set_u = find(u)
            set_v = find(v)
            if set_u == set_v:
                continue
            mst_edges.append(edge)
            union(set_u, set_v)
            num_components -= 1

    return sum(map(lambda x: x[2], mst_edges))


# All modules for CS 412 must include a main method that allows it
# to imported and invoked from other python scripts
def main():
    n: int = int(input())

    coords: CoordMap = {}

    for i in range(0, n):
        tokens = input().split()
        coords[i] = (float(tokens[0]), float(tokens[1]))

    weighted_graph = create_graph(coords)

    cost = calculate_cost(weighted_graph)

    print(f"${cost:.1f}M")


if __name__ == "__main__":
    main()