add method to print tree.

1 files changed, 25 insertions, 17 deletions

diff --git a/lib/tree.rb b/lib/tree.rb
index e4527d8..a03ff43 100644
--- a/lib/tree.rb
+++ b/lib/tree.rb
@@ -1,22 +1,30 @@
-  class Tree
-    attr_accessor :value, :left, :right
+class Tree
+  attr_accessor :value, :left, :right
 
-    def initialize(value, left: nil, right: nil)
-      @value = value
-      @left = left
-      @right = right
-    end
+  def initialize(value, left: nil, right: nil)
+    @value = value
+    @left = left
+    @right = right
+  end
+
+  def to_a
+    [value, left&.to_a, right&.to_a]
+  end
 
-    def print_tree
-      [value, left&.print_tree, right&.print_tree]
-    end
+  def print(prefix = '', tail = true)
+    puts(prefix + (tail ? "└── " : "├── ") + value.to_s)
 
-    def to_h
-      { value: value, left: left&.to_h, right: right&.to_h }
-    end
+    prefix = prefix + (tail ? "    " : "│   ")
+    left.print(prefix, false) if left
+    right.print(prefix, false) if right
+  end
+
+  def to_h
+    { value: value, left: left&.to_h, right: right&.to_h }
+  end
 
-    def self.build_from(hash)
-      return nil if hash.nil?
-      Tree.new(hash[:value], left: build_from(hash[:left]), right: build_from(hash[:right]))
-    end
+  def self.build_from(hash)
+    return nil if hash.nil?
+    Tree.new(hash[:value], left: build_from(hash[:left]), right: build_from(hash[:right]))
   end
+end