Flatten a tree (list of lists) with one statement?

.net tree flatten
Bob Tway picture Bob Tway · Oct 11, 2010 · Viewed 10.5k times · Source

Thanks to nHibernate, some of the data structures I work with are lists within lists within lists. So for example I have a data object called "category" which has a .Children property that resolves to a list of categories ... each one of which can have children ... and so on and so on.

I need to find a way of starting at a top-level category in this structure and getting a list or array or something similar of all the children in the entire structure - so all the children of all the children etc etc, flattened into a single list.

I'm sure it can be done with recursion, but I find recursive code a pain to work through, and I'm convinced there must be a more straightforward way in .Net 4 using Linq or somesuch - any suggestions?

Answer

Thomas Levesque picture Thomas Levesque · Oct 11, 2010

Here's an extension method that does the job:

// Depth-first traversal, recursive
public static IEnumerable<T> Flatten<T>(
        this IEnumerable<T> source,
        Func<T, IEnumerable<T>> childrenSelector)
{
    foreach (var item in source)
    {
        yield return item;
        foreach (var child in childrenSelector(item).Flatten(childrenSelector))
        {
            yield return child;
        }
    }
}

You can use it like this:

foreach(var category in categories.Flatten(c => c.Children))
{
    ...
}

The solution above does a depth-first traversal, if you want a breadth-first traversal you can do something like this:

// Breadth-first traversal, non-recursive
public static IEnumerable<T> Flatten2<T>(
        this IEnumerable<T> source,
        Func<T, IEnumerable<T>> childrenSelector)
{
    var queue = new Queue<T>(source);
    while (queue.Count > 0)
    {
        var item = queue.Dequeue();
        yield return item;
        foreach (var child in childrenSelector(item))
        {
            queue.Enqueue(child);
        }
    }
}

It also has the benefit of being non-recursive...

UPDATE: Actually, I just thought of a way to make the depth-first traversal non-recursive... here it is:

// Depth-first traversal, non-recursive
public static IEnumerable<T> Flatten3<T>(
        this IEnumerable<T> source,
        Func<T, IEnumerable<T>> childrenSelector)
{
    LinkedList<T> list = new LinkedList<T>(source);
    while (list.Count > 0)
    {
        var item = list.First.Value;
        yield return item;
        list.RemoveFirst();
        var node = list.First;
        foreach (var child in childrenSelector(item))
        {
            if (node != null)
                list.AddBefore(node, child);
            else
                list.AddLast(child);
        }
    }
}

I'm using a LinkedList<T> because insertions are O(1) operations, whereas insertions to a List<T> are O(n).

Related questions

How do I print out a tree structure?

I'm trying to improve performance in our app. I've got performance information in the form of a tree of calls, with the following node class: public class Node { public string Name; // method name public decimal Time; // time spent in method …

Read More
How to flatten tree via LINQ?

So I have simple tree: class MyNode { public MyNode Parent; public IEnumerable<MyNode> Elements; int group = 1; } I have a IEnumerable<MyNode>. I want to get a list of all MyNode (including inner node objects (Elements)) as …

Read More
Objects that represent trees

Are there any objects in C# (or in .net) that represents a binary tree (or for curiosity) and n-ary tree? I am not talking about presentation tree controls, but as model objects. If not, are there any good external implementations?

Read More