Search code examples
c#.netlinq.net-4.0tree

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 one flat list Where group == 1. How to do such thing via LINQ?


Solution

  • You can flatten a tree like this:

    IEnumerable<MyNode> Flatten(IEnumerable<MyNode> e) =>
        e.SelectMany(c => Flatten(c.Elements)).Concat(new[] { e });
    

    You can then filter by group using Where(...).

    To earn some "points for style", convert Flatten to an extension function in a static class.

    public static IEnumerable<MyNode> Flatten(this IEnumerable<MyNode> e) =>
        e.SelectMany(c => c.Elements.Flatten()).Concat(e);
    

    To earn more points for "even better style", convert Flatten to a generic extension method that takes a tree and a function that produces descendants from a node:

    public static IEnumerable<T> Flatten<T>(
        this IEnumerable<T> e
    ,   Func<T,IEnumerable<T>> f
    ) => e.SelectMany(c => f(c).Flatten(f)).Concat(e);
    

    Call this function like this:

    IEnumerable<MyNode> tree = ....
    var res = tree.Flatten(node => node.Elements);
    

    If you would prefer flattening in pre-order rather than in post-order, switch around the sides of the Concat(...).