How to flatten a deep tree using LINQ Flatten()

Question

I have a generic type tree class as follows:

public class TreeNode<T>
{
    private readonly T _data;
    private readonly TreeNode<T> _parent;
    private readonly List<TreeNode<T>> _children;
    private readonly int _level;

    public int Level
    { get { return _level; } }
    public int Count
    { get { return _children.Count; } }
    public bool IsRoot
    { get { return _parent == null; } }
    public bool IsLeaf
    { get { return _children.Count == 0; } }
    public T Data
    { get { return _data; } }
    public TreeNode<T> Parent
    { get { return _parent; } }

    public TreeNode(T data)
    {
        _data = data;
        _children = new List<TreeNode<T>>();
        _level = 0;
    }

    public TreeNode(T data, TreeNode<T> parent) : this(data)
    {
        _parent = parent;
        _level = _parent != null ? _parent.Level + 1 : 0;
    }

    public TreeNode<T> this[int key]
    {
        get { return _children[key]; }
    }

    public void AddChild(TreeNode<T> value)
    {
        _children.Add(value);
    }

    public TreeNode<T> GetChild(int key)
    {
        return _children[key];
    }

    public void Clear()
    {
        _children.Clear();
    }
}

I want to create a function in this class to flatten all the nested children list, no matter how deep the tree is. The return value would be a list of TreeNode. The function can be used on any node, no matter it is root or not.

I tried solution using LINQ Flatten() but it seems to return only the first level. How can I design the function in the TreeNode class using LINQ (or not) so that no matter what type the node is I can always return such list?

Answer 1

You can create enumerator function to enumerate all children of a node. Here is an example of how to do it:

public IEnumerable<TreeNode<T>> EnumerateChildrenRecursive()
{
    foreach (var child in _children)
    {
        yield return child;

        foreach (var subItem in child.EnumerateChildrenRecursive())
        {
            yield return subItem;
        }
    }
}

Converting to list will be simple:

var list = root.EnumerateChildrenRecursive().ToList();

Or if you have VERY deep tree we can emulate recursive algorithm:

public IEnumerable<TreeNode<T>> EnumerateChildrenFlatten()
{
    var stack = new Stack<TreeNode<T>>();

    foreach (var child in _children)
    {
        stack.Push(child);
    }

    while (stack.TryPop(out var current))
    {
        yield return current;

        foreach (var child in current._children)
        {
            stack.Push(child);
        }
    }
}

Answer 2

I'm not sure what this Flatten()-method is, it is not a built in method.

But you need to traverse the tree. My typical approach to this is a extension method following the same approach as LINQ:

public static IEnumerable<T> DepthFirst<T>(T self, Func<T, IEnumerable<T>> selector)
{
    var stack = new Stack<T>();
    stack.Push(self);
    while (stack.Count > 0)
    {
        var current = stack.Pop();
        yield return current;
        foreach (var child in selector(current))
        {
            stack.Push(child);
        }
    }
}

You will need to expose a IEnumerable of your children on your node class:

public class TreeNode<T>
{
   ...
   public IEnumerable<TreeNode<T>> Children => _children
}
...
var allNodes = myRootNode.DepthFirst(n => n.Children);

You can change the iteration order to breadth first by changing the stack to a Queue.

This requires that the tree is actually a tree and not a graph. If a TreeNode has itself as a child the returned IEnumerable will be infinitely long and you will likely get a stack overflow. To prevent this you can add list of already visited nodes to prevent visiting the same node multiple times.