A Ruby implementation of a Tree data structure and Tree Traversal algorithms

tree.rb

require 'test/unit'

module Tree
    class Node
        include Comparable
        include Enumerable

        attr_accessor :parent, :children, :element

        def initialize(element)
            @parent = nil
            @children = []
            @element = element
        end

        def root?
            @parent = nil
        end

        def leaf?
            @children.size == 0
        end

        alias_method :external?, :leaf?

        def left?
            !!@children[0]
        end

        def right?
            !!@children[1]
        end

        def left
            @children[0]
        end

        def right
            @children[1]
        end

        def sibling
            return nil if root?
            return parent.right if self == parent.left
            return parent.left if self == parent.right 
        end

        def [](index)
            @children[index]
        end

        def add_child_at(index, child)
            @children[index] = child
            child.parent = self if child
            child
        end

        def add(child)
            if child
                @children << child
                child.parent = self
            end

            child
        end

        alias_method :<<, :add

        def left=(child)
            add_child_at 0, child 
        end

        def right=(child)
            add_child_at 1, child
        end

        # The depth is a number of ancestors of a node p, excluding p itself
        # O(dp + 1) -> O(n)
        def depth
            return 0 if root?
            1 + parent.height
        end

        # The height of a nonempty tree is equal to the maximum of the depths of its leaf positions
        # O(n^2)
        def height
            return 0 if leaf?
            1 + children.map { |child| child.height }.max
        end

        def <=> (other)
            self.element <=> other.element
        end

        # O(n)
        def each_preorder(&block)
            yield self
            children.each { |child| child.each_preorder &block }

            self if block_given?
        end

        alias_method :each, :each_preorder

        # O(n) - symmetric 
        def each_inorder(&block)
            left.each_inorder &block if left?
            yield self
            right.each_inorder &block if right?

            self if block_given?
        end

        # O(n)
        def each_postorder(&block)
            children.each { |child| child.each_postorder &block }
            yield self

            self if block_given?
        end

        def each_breadthfirst(&block)
            queue = [self]
            until queue.empty?
                current = queue.shift
                yield current
                current.children.each { |child| queue << child }
            end 

            self if block_given?
        end

        alias_method :each_depthfirst, :each_postorder
    end
end

if __FILE__ == $0
    class TestTree < Test::Unit::TestCase
        def setup
            @root             = Tree::Node.new 0
            @root.left        = Tree::Node.new 1
            @root.right       = Tree::Node.new 2

            @root.left.left   = Tree::Node.new 3
            @root.left.right  = Tree::Node.new 4
            
            @root.right.left  = Tree::Node.new 5
            @root.right.right = Tree::Node.new 6
        end

        def teardown
            @root = nil
        end

        def test_preorder
            assert_equal [0, 1, 3, 4, 2, 5, 6], tree_map(:each_preorder)
        end

        def test_inorder
            assert_equal [3, 1, 4, 0, 5, 2, 6], tree_map(:each_inorder)
        end

        def test_postorder
            assert_equal [3, 4, 1, 5, 6, 2, 0], tree_map(:each_postorder)
        end

        def test_breadthfirst
            assert_equal [0, 1, 2, 3, 4, 5, 6], tree_map(:each_breadthfirst)
        end

    private

        def tree_map(method)
            result = []
            @root.send method do |node|
                result << node.element
            end
            result
        end
    end
end