prim.ts

import { PriorityQueue } from '../data_structures/heap/heap'
/**
 * @function prim
 * @description Compute a minimum spanning tree(MST) of a fully connected weighted undirected graph. The input graph is in adjacency list form. It is a multidimensional array of edges. graph[i] holds the edges for the i'th node. Each edge is a 2-tuple where the 0'th item is the destination node, and the 1'th item is the edge weight.
 * @Complexity_Analysis
 * Time complexity: O(Elog(V))
 * Space Complexity: O(V)
 * @param {[number, number][][]} graph - The graph in adjacency list form
 * @return {Edge[], number} - [The edges of the minimum spanning tree, the sum of the weights of the edges in the tree]
 * @see https://en.wikipedia.org/wiki/Prim%27s_algorithm
 */
export const prim = (graph: [number, number][][]): [Edge[], number] => {
  if (graph.length == 0) {
    return [[], 0];
  }
  let minimum_spanning_tree: Edge[] = [];
  let total_weight = 0;

  let priorityQueue = new PriorityQueue((e: Edge) => { return e.b }, graph.length, (a: Edge, b: Edge) => { return a.weight < b.weight });
  let visited = new Set<number>();

  // Start from the 0'th node. For fully connected graphs, we can start from any node and still produce the MST.
  visited.add(0);
  add_children(graph, priorityQueue, 0);

  while (!priorityQueue.isEmpty()) {
    // We have already visited vertex `edge.a`. If we have not visited `edge.b` yet, we add its outgoing edges to the PriorityQueue.
    let edge = priorityQueue.extract();
    if (visited.has(edge.b)) {
      continue;
    }
    minimum_spanning_tree.push(edge);
    total_weight += edge.weight;
    visited.add(edge.b);
    add_children(graph, priorityQueue, edge.b);
  }

  return [minimum_spanning_tree, total_weight];
}

const add_children = (graph: [number, number][][], priorityQueue: PriorityQueue<Edge>, node: number) => {
  for (let i = 0; i < graph[node].length; ++i) {
    let out_edge = graph[node][i];
    // By increasing the priority, we ensure we only add each vertex to the queue one time, and the queue will be at most size V.
    priorityQueue.increasePriority(out_edge[0], new Edge(node, out_edge[0], out_edge[1]));
  }
}

export class Edge {
  a: number = 0;
  b: number = 0;
  weight: number = 0;
  constructor(a: number, b: number, weight: number) {
    this.a = a;
    this.b = b;
    this.weight = weight;
  }
}