myWallJSON myWallJSON - 28 days ago 10
C# Question

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?

Answer

You can flatten a tree like this:

IEnumerable<MyNode> Flatten(IEnumerable<MyNode> e) {
    return 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) {
    return e.SelectMany(c => c.Elements.Flatten()).Concat(e);
}

To earn some points for "even better style", convert Flatten to a generic extension method that takes a tree and a function that produces descendents:

public static IEnumerable<T> Flatten<T>(
    this IEnumerable<T> e,
    Func<T,IEnumerable<T>> f) 
{
    return 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(...).