Как сделать IEnumerable значений узлов дерева?

Question

Рассмотрим следующий класс:

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; 
    } 
} 

Каждое значение передается O(h) раз, когда h высота дерева. Сначала в инструкции yield return value;, а затем в yield return i; каждого родительского узла.

Таким образом, sub_values возвращает n значениями с использованием O(nh) временной сложности.

Я могу заменить его на метод, который принимает ссылку на список целых чисел и вместо возвращаемых значений добавляет их в этот список, но он больше не будет лениться.

Могу ли я вернуть n значения в O(n) и поддерживать лень?

Answer 1

Это очень похоже на Function which will return particular node from tree structure и другие сообщения SO относительно рекурсивных итераторов. Все они включают использование явного стека или очереди. Вот общее решение для любого типа дерева. Пусть сначала определить многоразовую функцию в некотором общем месте, так что в следующий раз 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(); 
     } 
    } 
} 

Теперь давайте определим некоторые полезные хелперы внутри вашего TreeNode класса

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); 
} 

Обратите внимание на метод Traverse - это обеспечивает лени ты по просят. Теперь вы можете использовать обычные методы LINQ для фильтрации, прогнозов и т. Д. Например, этот метод становится таким, как это

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