How to make an IEnumerable of values of tree nodes?

Question

Consider the following class:

class TreeNode
{
    int value;
    public TreeNode l, r;
    public TreeNode(int value)
    {
        this.value = value;
    }
    public IEnumerable<int> sub_values()
    {
        if (l != null)
            foreach (int i in l.sub_values())
                yield return i;
        if (r != null)
            foreach (int i in r.sub_values())
                yield return i;
        yield return value;
    }
}

Each value is passed O(h) times, where h is the height of the tree. First, in yield return value; statement, then in yield return i; of each parent node.

So, sub_values returns n values using O(nh) time complexity.

I can replace it with a method, which accepts a reference to a list of integers and instead of returning values, adds them to this list, but it won't be lazy anymore.

Can I return n values in O(n) and maintain laziness?

Answer 1

It's very similar to Function which will return particular node from tree structure and other SO posts regarding recursive iterators. All they involve using an explicit stack or queue. Here is a generic solution for any type of tree. Let first define a reusable function in some common place, so the next time DRY

public static class TreeHelper
{
    public static IEnumerable<T> Traverse<T>(T node, Func<T, IEnumerable<T>> childrenSelector, bool preOrder = true)
    {
        var stack = new Stack<IEnumerator<T>>();
        var e = Enumerable.Repeat(node, 1).GetEnumerator();
        try
        {
            while (true)
            {
                while (e.MoveNext())
                {
                    var item = e.Current;
                    var children = childrenSelector(item);
                    if (children == null)
                        yield return item;
                    else
                    {
                        if (preOrder) yield return item;
                        stack.Push(e);
                        e = children.GetEnumerator();
                    }
                }
                if (stack.Count == 0) break;
                e.Dispose();
                e = stack.Pop();
                if (!preOrder) yield return e.Current;
            }
        }
        finally
        {
            e.Dispose();
            while (stack.Count != 0) stack.Pop().Dispose();
        }
    }
}

Now let define some useful helpers inside your TreeNode class

public bool AnyChildren() { return l != null || r != null; }
public IEnumerable<TreeNode> Children
{
    get
    {
        if (l != null) yield return l;
        if (r != null) yield return r;
    }
}
public IEnumerable<TreeNode> Traverse(bool preOrder = false)
{
    return TreeHelper.Traverse(this, node => node.AnyChildren() ? node.Children : null, preOrder);
}

Notice the Traverse method - it provides the laziness you are asking for. Now you can use the usual LINQ methods for filtering, projections etc. For instance, the method in question becomes like this

public IEnumerable<int> sub_values()
{
    return Traverse().Select(node => node.value);
}