Interview question about tree reconstruction.

traversal.rb

# Exercise:
#   Given the Pre-order and In-order traversals of a tree,
#   reconstruct the original tree structure.
#
#   Pre-order node output: Node (Left) (Right)
#   In-order node output:  (Left) Node (Right)
#
#   Example:
#                 'a'
#               |     |
#              'b'   'c'
#                   |   |
#                  'd' 'e'
#
#  In-order:  "badce"
#  Pre-order: "abcde
#
class Node
  # Build a tree structure via divide and conquer using subtree traversal outputs.

  def initialize(io, po)
    # Value of node is always first element of pre-order traversal
    @value = po.first

    # Elements in the IO traversal prior to @value belong to the left subtree
    @left = if (l_segment = io.take_while { |n| n != @value }).empty?
              nil
            else
              # After @value, the same number of elements occur from the left
              # subtree in PO traversal.
              self.class.new(l_segment, po.drop(1).take(l_segment.size))
            end

    # In the IO traversal, the right subtree is all elements following the left
    # subtree (size-known) and @value.
    @right =  if (r_segment = io.drop(l_segment.size + 1)).empty?
                nil
              else
                # Similarly, all remaining elements in the PO traversal belong
                # to the right subtree.
                self.class.new(r_segment, po.drop(l_segment.size + 1))
              end
  end

  # Methods for outputting the traversals.
  def in_order
    left = @left ? @left.in_order << " " : ""
    right = @right ? " " << @right.in_order : ""

    "#{left}#{@value}#{right}"
  end

  def pre_order
    left = @left ? " " << @left.pre_order : ""
    right = @right ? " " << @right.pre_order : ""

    "#{@value}#{left}#{right}"
  end

end

tree = Node.new(%w{b a d c e}, %w{a b c d e})
 # => #<Node:0x007fec1f2ba5d0
 #     @left=#<Node:0x007fec1f2ba508 @left=nil, @right=nil, @value="b">,
 #     @right=
 #      #<Node:0x007fec1f2ba440
 #       @left=#<Node:0x007fec1f2ba3a0 @left=nil, @right=nil, @value="d">,
 #       @right=#<Node:0x007fec1f2ba2d8 @left=nil, @right=nil, @value="e">,
 #       @value="c">,
 #     @value="a">

tree.in_order # => "b a d c e"
tree.pre_order # => "a b c d e"