Prim's Algorithm in Python, Java and TypeScript

import heapq


def prim(graph, start):
    in_tree = {start}
    tree = []
    heap = [(weight, start, neighbour) for neighbour, weight in graph[start]]
    heapq.heapify(heap)
    while heap and len(in_tree) < len(graph):
        weight, u, v = heapq.heappop(heap)
        if v in in_tree:
            continue
        in_tree.add(v)
        tree.append((u, v, weight))
        for neighbour, next_weight in graph[v]:
            if neighbour not in in_tree:
                heapq.heappush(heap, (next_weight, v, neighbour))
    return tree


def undirected(edges):
    graph = {}
    for u, v, weight in edges:
        graph.setdefault(u, []).append((v, weight))
        graph.setdefault(v, []).append((u, weight))
    return graph


EDGES = [
    ("A", "B", 4), ("A", "C", 1), ("B", "C", 2), ("B", "D", 5), ("C", "D", 8),
    ("C", "E", 9), ("D", "E", 3), ("D", "F", 7), ("E", "F", 6),
]

tree = prim(undirected(EDGES), "A")
print("added:", ", ".join(f"{u}-{v} {weight}" for u, v, weight in tree))
print("total:", sum(weight for _, _, weight in tree))

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.stream.Collectors;

public class PrimCode {
    record Edge(String u, String v, int weight) {
        @Override
        public String toString() {
            return u + "-" + v + " " + weight;
        }
    }

    static List<Edge> prim(Map<String, List<Edge>> graph, String start) {
        Set<String> inTree = new HashSet<>(Set.of(start));
        List<Edge> tree = new ArrayList<>();
        PriorityQueue<Edge> heap = new PriorityQueue<>(Comparator.comparingInt(Edge::weight));
        heap.addAll(graph.get(start));
        while (!heap.isEmpty() && inTree.size() < graph.size()) {
            Edge edge = heap.poll();
            if (!inTree.add(edge.v())) continue;
            tree.add(edge);
            for (Edge next : graph.get(edge.v())) {
                if (!inTree.contains(next.v())) heap.add(next);
            }
        }
        return tree;
    }

    static Map<String, List<Edge>> undirected(List<Edge> edges) {
        Map<String, List<Edge>> graph = new LinkedHashMap<>();
        for (Edge edge : edges) {
            graph.computeIfAbsent(edge.u(), key -> new ArrayList<>()).add(edge);
            graph.computeIfAbsent(edge.v(), key -> new ArrayList<>()).add(new Edge(edge.v(), edge.u(), edge.weight()));
        }
        return graph;
    }

    public static void main(String[] args) {
        List<Edge> edges = List.of(
            new Edge("A", "B", 4), new Edge("A", "C", 1), new Edge("B", "C", 2), new Edge("B", "D", 5), new Edge("C", "D", 8),
            new Edge("C", "E", 9), new Edge("D", "E", 3), new Edge("D", "F", 7), new Edge("E", "F", 6));
        List<Edge> tree = prim(undirected(edges), "A");
        System.out.println("added: " + tree.stream().map(Edge::toString).collect(Collectors.joining(", ")));
        System.out.println("total: " + tree.stream().mapToInt(Edge::weight).sum());
    }
}

type Edge = [string, string, number];
type Neighbour = [string, number];

function prim(graph: Map<string, Neighbour[]>, start: string): Edge[] {
  const inTree = new Set<string>([start]);
  const best = new Map<string, [number, string]>();
  const tree: Edge[] = [];
  let latest = start;
  while (inTree.size < graph.size) {
    for (const [neighbour, weight] of graph.get(latest)!) {
      if (!inTree.has(neighbour) && weight < (best.get(neighbour)?.[0] ?? Infinity)) best.set(neighbour, [weight, latest]);
    }
    let next: string | null = null;
    for (const [vertex, [weight]] of best) {
      if (next === null || weight < best.get(next)![0]) next = vertex;
    }
    if (next === null) break;
    const [weight, from] = best.get(next)!;
    best.delete(next);
    inTree.add(next);
    tree.push([from, next, weight]);
    latest = next;
  }
  return tree;
}

function undirected(edges: Edge[]): Map<string, Neighbour[]> {
  const graph = new Map<string, Neighbour[]>();
  for (const [u, v, weight] of edges) {
    if (!graph.has(u)) graph.set(u, []);
    if (!graph.has(v)) graph.set(v, []);
    graph.get(u)!.push([v, weight]);
    graph.get(v)!.push([u, weight]);
  }
  return graph;
}

const edges: Edge[] = [
  ["A", "B", 4], ["A", "C", 1], ["B", "C", 2], ["B", "D", 5], ["C", "D", 8],
  ["C", "E", 9], ["D", "E", 3], ["D", "F", 7], ["E", "F", 6],
];
const tree = prim(undirected(edges), "A");
console.log(`added: ${tree.map(([u, v, weight]) => `${u}-${v} ${weight}`).join(", ")}`);
console.log(`total: ${tree.reduce((sum, [, , weight]) => sum + weight, 0)}`);