Cornel’s Blog

Implementing an Iterator for Binary Trees

Ξ August 5th, 2009 | → 0 Comments | ∇ .Net Framework, Algorithms, Tools and technologies |

Some time ago I wrote a post showing how to implement a generic binary search tree data structure in C#. In today’s post I will provide a way of extending the binary data structure in order to be able to use LINQ in order to query it. In order to be able to use LINQ over a custom data structure, all that we have to extend our data structure such that it implements the IEnumerable<T> interface. This is due to the fact that LINQ provides a set of extension methods that extend the functionality of types inheriting from IEnumerable<T>.

In order to support LINQ to Objects, the binary tree structure presented in this post was extended, by implementing the IEnumerable<T> interface. The Enumerator starts from the root of the tree and uses a depth first search approach in iterating through the nodes of the tree. The code bellow shows the implementation of the methods:

namespace TreeUtility
{
    using System;
    using System.Collections;
    using System.Collections.Generic;

    /// <summary>
    /// Represents a binary search tree
    /// </summary>
    /// <typeparam name="T">Specifies the element type
    /// of the binary search tree</typeparam>
    public partial class BinarySearchTree<T> :
        IEnumerable<Node<T>> where T : IComparable
    {
        #region IEnumerable<T> Members

        // returns an enumerator for the BinarySearchTree
        // collection
        public IEnumerator<Node<T>> GetEnumerator()
        {
            // we check if the collection contains elements
            if (root == null)
                throw new InvalidOperationException(
                    "Cannot iterate an empty collection");

            // variable use to track the current node
            Node<T> current;

            // structure used for keeping the stack of
            // discovered elements
            Stack<Node<T>> stack = new Stack<Node<T>>();

            // we use depth first search...
            // we start by pushing the root on the stack
            stack.Push(root);

            // as long as we still have elements in the
            // collection which were not yet iterated, we loop
            while (stack.Count > 0)
            {
                // the current element is the element from the
                // top of the stack
                current = stack.Pop();

                // we add it's children on the stack
                if (current.Left != null)
                {
                    stack.Push(current.Left);
                }

                if (current.Right != null)
                {
                    stack.Push(current.Right);
                }

                // and we return the element to the iterator
                yield return current;
            }
        }

        #endregion

        #region IEnumerable Members

        // method from the IEnumerable interface
        IEnumerator IEnumerable.GetEnumerator()
        {
            // just return the generic iterator
            return this.GetEnumerator();
        }

        #endregion

        // rest of the code ...
    }
}

The code to test the class was extended in order to test the iterator methods which are called when looping over the structure. The modified code is presented bellow:

namespace TreeUtility
{
    using System;
    using System.Linq;

    class Program
    {
        static void Main(string[] args)
        {
            Test();
            Console.ReadKey(true);
        }

        /// <summary>
        /// Tests the functionality of the BST
        /// </summary>
        static void Test()
        {
            // Creates an initial array
            int[] numbers = { 8, 1, 5, 7, 2, 3, 9, 10, 2, 4, 6, 12 };

            // Creates a new instance for of a BST
            BinarySearchTree<int> bst = new BinarySearchTree<int>();

            // Adds some new elements in the tree
            foreach (int number in numbers)
            {
                bst.Add(number);
            }

            // removes an element from the tree
            bst.Remove(1);

            // use the iterator to traverse the collection
            foreach (var item in bst)
            {
                Console.WriteLine(item.Info);
            }

            // Displays the trees structure
            Console.WriteLine(bst.ToString());

            // use LINQ to filter the collection elements and
            // retrieve only the odd nodes
            var oddNodes = from x in bst
                           where x.Info % 2 == 1
                           select x.Info;

            // display the results of the LINQ query
            foreach (var oddNode in oddNodes)
            {
                Console.WriteLine(oddNode);
            }

        }
    }
}

In conclusion, with a minimum amount of effort, by implementing the IEnumerable<T> interface, one can fully leverage the functionality provided by LINQ in order to query and filter custom data structures.

Leave a reply