Skip to content

Commit

Permalink
feat: add Prim's algorithm for Minimum Spanning Tree
Browse files Browse the repository at this point in the history
fix: do not decrease priority in PriorityQueue increasePriority()
  • Loading branch information
caojoshua committed Jul 6, 2023
1 parent 8deeca2 commit 1507c47
Show file tree
Hide file tree
Showing 4 changed files with 154 additions and 1 deletion.
6 changes: 5 additions & 1 deletion data_structures/heap/heap.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
export abstract class Heap<T> {
protected heap: T[];
// A comparison function. Returns true if a should be the parent of b.
private compare: (a: T, b: T) => boolean;
protected compare: (a: T, b: T) => boolean;

constructor(compare: (a: T, b: T) => boolean) {
this.heap = [];
Expand Down Expand Up @@ -189,6 +189,10 @@ export class PriorityQueue<T> extends MinHeap<T> {
return;
}
let key = this.keys[idx];
if (this.compare(this.heap[key], value)) {
// Do not do anything if the value in the heap already has a higher priority.
return;
}
// Increase the priority and bubble it up the heap.
this.heap[key] = value;
this.bubbleUp(key);
Expand Down
5 changes: 5 additions & 0 deletions data_structures/heap/test/min_heap.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,11 @@ describe("MinHeap", () => {
heap.increasePriority(81, 72);
heap.increasePriority(9, 0);
heap.increasePriority(43, 33);
// decreasing priority should do nothing
heap.increasePriority(72, 100);
heap.increasePriority(12, 24);
heap.increasePriority(39, 40);

heap.check();
// Elements after increasing priority
const newElements: number[] = [
Expand Down
59 changes: 59 additions & 0 deletions graph/prim.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
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 always store the previously visited edge in `edge.a`, and the newly visited node in `edge.b`.
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;
}
}

85 changes: 85 additions & 0 deletions graph/test/prim.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
import { Edge, prim } from "../prim";

let edge_equal = (x: Edge, y: Edge): boolean => {
return (x.a == y.a && x.b == y.b) || (x.a == y.b && x.b == y.a) && x.weight == y.weight;
}

let test_graph = (expected_tree_edges: Edge[], other_edges: Edge[], num_vertices: number, expected_cost: number) => {
// First make sure the graph is undirected
let graph: [number, number][][] = [];
for (let _ = 0; _ < num_vertices; ++_) {
graph.push([]);
}
for (let edge of expected_tree_edges) {
graph[edge.a].push([edge.b, edge.weight]);
graph[edge.b].push([edge.a, edge.weight]);
}
for (let edge of other_edges) {
graph[edge.a].push([edge.b, edge.weight]);
graph[edge.b].push([edge.a, edge.weight]);
}

let [tree_edges, cost] = prim(graph);
expect(cost).toStrictEqual(expected_cost);
for (let expected_edge of expected_tree_edges) {
expect(tree_edges.find(edge => edge_equal(edge, expected_edge))).toBeTruthy();
}
for (let unexpected_edge of other_edges) {
expect(tree_edges.find(edge => edge_equal(edge, unexpected_edge))).toBeFalsy();
}
};


describe("prim", () => {

it("should return empty tree for empty graph", () => {
expect(prim([])).toStrictEqual([[], 0]);
});

it("should return empty tree for single element graph", () => {
expect(prim([])).toStrictEqual([[], 0]);
});

it("should return correct value for two element graph", () => {
expect(prim([[[1, 5]], []])).toStrictEqual([[new Edge(0, 1, 5)], 5]);
});

it("should return the correct value", () => {
let expected_tree_edges = [
new Edge(0, 1, 1),
new Edge(1, 3, 2),
new Edge(3, 2, 3),
];

let other_edges = [
new Edge(0, 2, 4),
new Edge(0, 3, 5),
new Edge(1, 2, 6),
];

test_graph(expected_tree_edges, other_edges, 4, 6);
});

it("should return the correct value", () => {
let expected_tree_edges = [
new Edge(0, 2, 2),
new Edge(1, 3, 9),
new Edge(2, 6, 74),
new Edge(2, 7, 8),
new Edge(3, 4, 3),
new Edge(4, 9, 9),
new Edge(5, 7, 5),
new Edge(7, 9, 4),
new Edge(8, 9, 2),
]

let other_edges = [
new Edge(0, 1, 10),
new Edge(2, 4, 47),
new Edge(4, 5, 42),
];

test_graph(expected_tree_edges, other_edges, 10, 116);
});

})

0 comments on commit 1507c47

Please sign in to comment.