Tree traversal, LINQ style

When dealing with tree structures, it’s helpful to be able to visit all nodes. Here’s an extension method that makes traversing tree structures easier (in this case, in preorder, depth-first).

internal static class TreeTraverse
{
    public static IEnumerable<T> PreorderTraverse<T>(this T node, Func<T, IEnumerable<T>> childrenFor)
    {
        yield return node;

        var childNodes = childrenFor(node);
        if (childNodes != null)
        {
            foreach (var childNode in childNodes.SelectMany(n => n.PreorderTraverse(childrenFor)))
            {
                yield return childNode;
            }
        }
    }
}

This extension method allows you to treat any type as the root node of a tree, providing you can define a sensible childrenFor method for it.

Example usage:

var root = new Node
{
    Name = "root",
    Children = new List<Node>()
    {
        new Node
        {
            Name = "root > child A",
            Children = new List<Node>()
            {
                new Node
                {
                    Name = "root > child A > child AA"
                }
            }
        },
        new Node
        {
            Name = "root > child B"
        }
    }
};

foreach (var node in root.PreorderTraverse(n => n.Children))
{
    Console.WriteLine(node.Name);
}

Where Node is defined as:

class Node
{
    public string Name { get; set; }
    public IEnumerable<Node> Children { get; set; }
}

To use PreorderTraverse on an IEnumerable<T> you can combine with SelectMany, as in fact the PreorderTraverse extension method does.

Comments are closed.