Oct 25 2008

Walking a tree with yield return and recursion

Category: Matt @ 06:32

C# 2.0 introduced the "yield return" statement.  While neat in theory, I actually never ran into a scenario where I needed it.  I had heard all the rage about it, but to me, it was kinda like the 5th wheel on the cart of awesomeness that was C# 2.0.  My opinion changed this morning when I needed a simple way to depth-first iterate over the items in a tree.

My original solution was to use a stack, build up a collection of the nodes in the tree, then return the collection, but that actually didn't work too well in my case (I was doing something screwy with the collection later on that was making ActiveRecord angry).  It also wasted memory.  Yeah, memory is cheap, but why waste it?

After thinking on it for a little bit, I decided I'd implement a custom IEnumerator for my tree, but then I remembered that I could just use 'yield return'.  5 lines of code later, we have it:

   1: public IEnumerator<TreeNode> GetDepthFirstEnumerator()
   2: {
   3:     yield return this;
   4:     
   5:     foreach (TreeNode child in ChildNodes)
   6:     {
   7:         IEnumerator<TreeNode> childEnumerator = child.GetDepthFirstEnumerator();
   8:  
   9:         while (childEnumerator.MoveNext())
  10:         {
  11:             yield return childEnumerator.Current;
  12:         }
  13:     }
  14: }

Surprisingly, this works fine even with recursion!

Tags:

blog comments powered by Disqus