Skip to content

Instantly share code, notes, and snippets.

@uztbt

uztbt/PrimsMST.ts

Created June 20, 2021 13:29
Download ZIP

Revisions

  1. uztbt created this gist Jun 20, 2021.
    55 changes: 55 additions & 0 deletions PrimsMST.ts
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,55 @@
    function prims(n: number, edges: number[][], start: number): number {
    // Write your code here
    const edgeDict: {[s: number]: number[][]} = {};
    // Create the edgeDict
    for (const edge of edges) {
    const [s, t] = [edge[0], edge[1]];
    if (typeof edgeDict[s] === "undefined") edgeDict[s] = [edge];
    else edgeDict[s].push(edge);
    if (typeof edgeDict[t] === "undefined") edgeDict[t] = [edge];
    else edgeDict[t].push(edge);
    }
    const visited: Set<number> = new Set([start]);
    const candEdges: number[][] = edgeDict[start].sort((a, b) => b[2] - a[2]); // desc.
    let acc = 0;
    while (visited.size < n) {
    console.log(`visited = ${visited}`);
    while (candEdges.length > 0) {
    const e = candEdges.pop();
    if (visited.has(e[0]) && !visited.has(e[1])) {
    console.log(`chose ${e}`);
    acc += e[2];
    visited.add(e[1]);
    const addingEdges = edgeDict[e[1]];
    for (const addingEdge of addingEdges) {
    const insertingIndex = reverseBinarySearch(candEdges, addingEdge);
    candEdges.splice(insertingIndex, 0, addingEdge);
    }
    } else if (visited.has(e[1]) && !visited.has(e[0])) {
    console.log(`chose ${e}`);
    acc += e[2];
    visited.add(e[0]);
    const addingEdges = edgeDict[e[0]];
    for (const addingEdge of addingEdges) {
    const insertingIndex = reverseBinarySearch(candEdges, addingEdge);
    candEdges.splice(insertingIndex, 0, addingEdge);
    }
    }
    // visited.has(e[0]) && visited.has(e[0]) ==> discard the edge
    }
    }
    return acc;
    }

    function reverseBinarySearch(arr: number[][], edge: number[]): number {
    let [l, r] = [0, arr.length];
    while(l < r) {
    const m = Math.floor((l+r)/2);
    if (arr[m][2] > edge[2]) {
    l = m + 1;
    } else {
    r = m;
    }
    }
    return l;
    }