Monday, April 1, 2013

Prim's Minimum Spanning Tree

Taking a break from Sasa blog posts, I ran into a problem that would be nicely solved by extracting the minimum spanning tree. Turns out this is pretty simple to solve naively in C# using only IEnumerable<T>:

/// <summary>
/// Computes the minimum spanning tree.
/// </summary>
/// <param name="edges">The list of edges.</param>
/// <remarks>
/// Given a list of edges whose vertices are consecutive numbers starting
/// from 0, this extension method returns an enumerable sequence of the
/// edges describing the minimum spanning tree.
/// </remarks>
public static IEnumerable<Edge> MinSpanTree(this IEnumerable<Edge> edges)
{
    var span = edges.ToList();
    var count = 1 + span.Max(x => Math.Max(x.From, x.To));
    edges.Sort();
    var vnew = new BitArray(count);
        vnew[0] = true;
    for (var i = 0; i < count; ++i)
    {
        foreach (var x in edges)
        {
            if (vnew[x.From] ^ vnew[x.To])
            {
                yield return x;
                vnew[x.From] = vnew[x.To] = true;
                break;
            }
        }
    }
}

An Edge is a simple class that encapsulates two vertex numbers and a weight represented by a double. It's also IComparable, so we can sort on it using the weight.

Prim's algorithm selects the minimum weighted edge between a vertex we have not yet accepted, and a vertex that we have accepted. We use a bit array to track the vertex numbers we've accepted, and simply search weight-sorted edge list for the first entry where one vertex is accepted, and the other is not. This is an O(V2), where V is the number of vertices. We could also contract the list after each accepted edge, but that doesn't change the overall time complexity.

A simple adjacency matrix representation using a two-dimensional array of doubles is also easily definable:

/// <summary>
/// Computes the minimum spanning tree.
/// </summary>
/// <param name="edges"></param>
/// <remarks>
/// Given an adjacency matrix describing the weighted edges, where
/// vertices are the indices of the matrix, this extension method
/// returns an enumerable sequence of the edges describing the
/// minimum spanning tree.
/// 
/// Vertices that have no edges between them should have
/// <see cref="double.NaN"/> or <see
/// cref="double.PositiveInfinity"/> in the corresponding
/// matrix entry.
/// </remarks>
public static IEnumerable<Edge> MinSpanTree(this double[,] edges)
{
    if (edges.RowCount() != edges.ColumnCount())
        throw new ArgumentException("Matrix row count must equal column count.", "edges");
    // convert the adjacency matrix into a list of edges
    var count = edges.RowCount();
    var span = new List<Edge>(count);
    for (int i = count - 1; i >= 0; --i)
    {
        for (int j = 0; j < i; ++j)
        {
            var w = edges[i, j];
            if (!double.IsNaN(w) && w < double.PositiveInfinity)
                span.Add(new Edge(i, j, w));
        }
    }
    return span.MinimumSpanningTree();
}

I've included these implementations in Sasa.Numerics.dll for the v0.9.4 release.

No comments: