How to flatten tree via LINQ?

Question

So I have simple tree:

class MyNode
{
 public MyNode Parent;
 public IEnumerable Elements;
 int group = 1;
}

I have a I

Answer 1

The problem with the accepted answer is that it is inefficient if the tree is deep. If the tree is very deep then it blows the stack. You can solve the problem by using an explicit stack:

public static IEnumerable<MyNode> Traverse(this MyNode root)
{
    var stack = new Stack<MyNode>();
    stack.Push(root);
    while(stack.Count > 0)
    {
        var current = stack.Pop();
        yield return current;
        foreach(var child in current.Elements)
            stack.Push(child);
    }
}

Assuming n nodes in a tree of height h and a branching factor considerably less than n, this method is O(1) in stack space, O(h) in heap space and O(n) in time. The other algorithm given is O(h) in stack, O(1) in heap and O(nh) in time. If the branching factor is small compared to n then h is between O(lg n) and O(n), which illustrates that the naïve algorithm can use a dangerous amount of stack and a large amount of time if h is close to n.

Now that we have a traversal, your query is straightforward:

root.Traverse().Where(item=>item.group == 1);

Answer 2

Just for completeness, here is the combination of the answers from dasblinkenlight and Eric Lippert. Unit tested and everything. :-)

 public static IEnumerable<T> Flatten<T>(
        this IEnumerable<T> items,
        Func<T, IEnumerable<T>> getChildren)
 {
     var stack = new Stack<T>();
     foreach(var item in items)
         stack.Push(item);

     while(stack.Count > 0)
     {
         var current = stack.Pop();
         yield return current;

         var children = getChildren(current);
         if (children == null) continue;

         foreach (var child in children) 
            stack.Push(child);
     }
 }

Answer 3

void Main()
{
    var allNodes = GetTreeNodes().Flatten(x => x.Elements);

    allNodes.Dump();
}

public static class ExtensionMethods
{
    public static IEnumerable<T> Flatten<T>(this IEnumerable<T> source, Func<T, IEnumerable<T>> childrenSelector = null)
    {
        if (source == null)
        {
            return new List<T>();
        }

        var list = source;

        if (childrenSelector != null)
        {
            foreach (var item in source)
            {
                list = list.Concat(childrenSelector(item).Flatten(childrenSelector));
            }
        }

        return list;
    }
}

IEnumerable<MyNode> GetTreeNodes() {
    return new[] { 
        new MyNode { Elements = new[] { new MyNode() }},
        new MyNode { Elements = new[] { new MyNode(), new MyNode(), new MyNode() }}
    };
}

class MyNode
{
    public MyNode Parent;
    public IEnumerable<MyNode> Elements;
    int group = 1;
}

Answer 4

In case anyone else finds this, but also needs to know the level after they've flattened the tree, this expands on Konamiman's combination of dasblinkenlight and Eric Lippert's solutions:

    public static IEnumerable<Tuple<T, int>> FlattenWithLevel<T>(
            this IEnumerable<T> items,
            Func<T, IEnumerable<T>> getChilds)
    {
        var stack = new Stack<Tuple<T, int>>();
        foreach (var item in items)
            stack.Push(new Tuple<T, int>(item, 1));

        while (stack.Count > 0)
        {
            var current = stack.Pop();
            yield return current;
            foreach (var child in getChilds(current.Item1))
                stack.Push(new Tuple<T, int>(child, current.Item2 + 1));
        }
    }

Answer 5

Here some ready to use implementation using Queue and returning the Flatten tree me first and then my children.

public static IEnumerable<T> Flatten<T>(this IEnumerable<T> items, 
    Func<T,IEnumerable<T>> getChildren)
    {
        if (items == null)
            yield break;

        var queue = new Queue<T>();

        foreach (var item in items) {
            if (item == null)
                continue;

            queue.Enqueue(item);

            while (queue.Count > 0) {
                var current = queue.Dequeue();
                yield return current;

                if (current == null)
                    continue;

                var children = getChildren(current);
                if (children == null)
                    continue;

                foreach (var child in children)
                    queue.Enqueue(child);
            }
        }

    }

Answer 6

A really other option is to have a proper OO design.

e.g. ask the MyNode to return all flatten.

Like this:

class MyNode
{
    public MyNode Parent;
    public IEnumerable<MyNode> Elements;
    int group = 1;

    public IEnumerable<MyNode> GetAllNodes()
    {
        if (Elements == null)
        {
            return Enumerable.Empty<MyNode>(); 
        }

        return Elements.SelectMany(e => e.GetAllNodes());
    }
}

Now you could ask the top level MyNode to get all the nodes.

var flatten = topNode.GetAllNodes();

If you can't edit the class, then this isn't an option. But otherwise, I think this is could be preferred of a separate (recursive) LINQ method.

This is using LINQ, So I think this answer is applicable here ;)