PrimsMST.ts

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;
}