math/MST.js

import UnionFind from "union-find";
import Delaunator from "delaunator";
import { point } from "./Utils.js";

/**
 * Minimum Spanning Tree implementation using union-find and kruskal
 * @example
 * 
 * //known points, no edges
 * let mst = new MST([[0, 0], [23, 34], [5, 645]]);
 * mst.edgesDelaunay();
 * let tree = msg.calculate();
 * 
 * //known points and known edges
 * let mst = new MST([[0, 0], [23, 34], [5, 645]]);
 * mst.edges = [[0,1], [1,2]];
 * let tree = msg.calculate();
 */
class MST {

    /**
     * @param {Array} vertices array of [points](module-Utils#~point)
     * @param {Function|null} measure measure function, default is manhattan distance
     */
    constructor(vertices, measure = null) {
        this.vertices = vertices.map(point);
        this.edges = [];
        this.measure = measure;

        if (this.measure === null) {
            //default measure is manhattan distance
            this.measure = (a, b) => {
                var dx = a[0] - b[0];
                var dy = a[1] - b[1];
                return dx * dx + dy * dy;
            };
        }

        this.tree = [];
    }

    /**
     * fill the edges array by connecting all edges
     */
    edgesDelaunay() {
        const delaunay = Delaunator.from(this.vertices);

        let edges = {};

        function addEdge(a, b) {
            let a_ = Math.min(a, b);
            let b_ = Math.max(a, b);
            edges[a_ + "|" + b_] = [a_, b_];
        }

        for (let i = 0; i < delaunay.triangles.length; i += 3) {
            const v1 = delaunay.triangles[i];
            const v2 = delaunay.triangles[i + 1];
            const v3 = delaunay.triangles[i + 2];

            addEdge(v1, v2);
            addEdge(v1, v3);
            addEdge(v2, v3);
        }

        this.edges = Object.values(edges);
    }

    /**
     * calculate minimum spanning tree using kruskal algorithm
     * @returns {Array}
     */
    calculate() {
        var finalEdge = [];

        var forest = new UnionFind(this.vertices.length);

        var edgeDist = [];
        for (var ind in this.edges) {
            var u = this.edges[ind][0];
            var v = this.edges[ind][1];
            var e = { edge: this.edges[ind], weight: this.measure(this.vertices[u], this.vertices[v]) };
            edgeDist.push(e);
        }

        edgeDist.sort(function(a, b) { return a.weight - b.weight; });

        for (var i = 0; i < edgeDist.length; i++) {
            var u = edgeDist[i].edge[0];
            var v = edgeDist[i].edge[1];

            if (forest.find(u) != forest.find(v)) {
                finalEdge.push([u, v]);
                forest.link(u, v);
            }
        }

        this.tree = finalEdge;
        return this.tree;
    }
}

export default MST