A js function for a chromatic spanning tree: a spanning tree that preserves the connection of same-color subgraphs.

trees.js

/* trees.js
 *
 * A couple spanning-tree algorithms.
 *
 * The most interesting function is the chromatic spanning tree,
 * which is intuitively simple yet (for me at least) was not obvious
 * in terms of finding a simple js implementation.
 *
 */

// This accepts a graph of the form {nodeId: nborsArray} and returns a spanning
// tree in the form [index1, kids1, index2, kids2, ... ], where each kids array
// is a list of the children of the previous index.
// This form of the spanning tree is designed to be easy to traverse in
// javascript.
export function findSpanningTree(graph) {
    // We'll perform a breadth-first search on the tree.

    let root    = parseInt(Object.keys(graph)[0]);
    let seen    = {[root]: true};
    let toVisit = [root];
    let tree    = [];

    while (toVisit.length > 0) {
        let node = toVisit.shift();
        let kids = [];
        for (let nbor of graph[node]) {
            if (nbor in seen) continue;
            kids.push(nbor);
            seen[nbor] = true;
            toVisit.push(nbor);
        }
        if (kids.length > 0) {
            tree.push(node);
            tree.push(kids);
        }
    }

    return tree;
}

// This works just like findSpanningTree(), except that it also
// preserves the connectedness of same-colored subgraphs. This expects
// the `colors` array to have the same length as the number of nodes in the
// graph, and that each same-color induced subgraph of `graph` is connected.
// The return value has the same format as in findSpanningTree().
export function findChromaticSpanningTree(graph, colors) {

    console.assert(Object.values(graph).length === colors.length);

    // We'll perform a breadth-first search on the tree.

    let root      = parseInt(Object.keys(graph)[0]);
    let seen      = {[root]: true};
    let toVisit   = [root];
    let tree      = [];
    let colorSrcs = {[colors[root]]: root};
    let newColors = [colors[root]];

    while (newColors.length > 0) {

        // Start a new color.
        let color = newColors.shift();
        toVisit = [colorSrcs[color]];

        while (toVisit.length > 0) {
            let node = toVisit.shift();
            let kids = [];
            for (let nbor of graph[node]) {
                if (nbor in seen) continue;
                let isNewColor = !(colors[nbor] in colorSrcs);
                if (colors[nbor] !== color && !isNewColor) continue;
                kids.push(nbor);
                seen[nbor] = true;
                if (isNewColor) {
                    colorSrcs[colors[nbor]] = nbor;
                    newColors.push(colors[nbor]);
                } else {
                    toVisit.push(nbor);
                }
            }
            if (kids.length > 0) {
                tree.push(node);
                tree.push(kids);
            }
        }
    }

    return tree;
}