二分木のAA

tree.rb

# coding: utf-8

# See http://ja.stackoverflow.com/questions/4739
#     https://gist.github.com/snipsnipsnip/0e267f6a39cc2397da3d
#
# 2 分木をコンソールにアスキーアートで出力する。
#

require 'gviz'  # gem install gviz

class Node
  attr_accessor :value, :left, :right
  def initialize(val)
    @value = val
    @left = nil
    @right = nil
  end
end

class Canvas
  def initialize(width, height)
    @canvas = Array.new(width){ Array.new(height, ' ') }  # 描画キャンバス
    @w, @h = [width, height]                              # canvas の幅, 高さ
  end

  # canvas の大きさを拡大させる。
  def expand(inc_x, inc_y)
    inc_x = 0 if inc_x < 0
    inc_y = 0 if inc_y < 0
    new_canvas = Array.new(@w + inc_x) { Array.new(@h + inc_y, ' ') }
    (0...@w).each do |x|
      (0...@h).each do |y|
        new_canvas[x][y] = @canvas[x][y]
      end
    end
    @canvas = new_canvas
    @w, @h = [@canvas.size, @canvas[0].size]
  end

  # canvas を STDOUT に表示する。
  def show
    (0...@h).each { |y| STDOUT.puts "|#{line(y)}|\n" }
  end

  def put_string(str, x, y)
    # 描画がはみ出るようなら、事前に canvas の大きさを拡張する。
    if (x + str.length >= @w) || y >= @h
      inc_x, inc_y = [0, 0]
      inc_x = 1 + (x + str.length - @w) if (x + str.length) >= @w
      inc_y = 1 + y - @h if y >= @h
      expand(inc_x, inc_y)
    end

    pos = x
    str.each_char do |c|
      @canvas[pos][y] = c
      pos += 1
    end
  end

  # 描画をする際に、既存の文字列を上書きしてしまわないかを調べる。
  def put_string?(str, x, y)
    return true if y < 0
    return true if y >= @h
    return true if x >= @w

    (0...str.length).each do |dx|
      return true if x + dx >= @w
      return false if (x + dx >= 0) && @canvas[x + dx][y] != ' '
    end
    true
  end

  def copy_area(x0, y0, dx, dy, fill = nil)
    strs = Array.new(dy, ' ')
    idx = 0
    (0...dy).each do |y|
      line = ''
      (0...dx).each do |x|
        line += "#{@canvas[x0 + x][y0 + y]}"
        @canvas[x0 + x][y0 + y] = fill if fill
      end
      strs[idx] = line
      idx += 1
    end
    strs
  end

  def clear(x0 = 0, y0 = 0, dx = @w, dy = @h)
    copy_area(x0, y0, dx, dy, ' ')
  end

  def fill_rect(c, x0, y0, dx, dy)
    copy_area(x0, y0, dx, dy, c)
  end

  def line(y)
    str = ''
    return '' if y < 0 || y >= @h
    (0...@w).each do |x|
      str += @canvas[x][y]
    end
    str
  end
end

# ツリーの宣言を短く書くために定義
def Node(val, left = nil, right = nil)
  node = Node.new(val)
  node.left = left
  node.right = right
  node
end

# ノードを描画できるかを調べる。(既存描画文字を壊さないかを知らべる)
def put_node?(canvas, str, x, y)
  return true if y < 0
  return false if x < 1
  return false unless canvas.put_string?('*' * (str.length + 1), x - 1, y)

  line_right = canvas.line(y)[x..-1]
  if line_right
    line_right.each_char do |c|
      return false if c != ' '
    end
  end
  true
end

# 2 分木を描画する。
def print_tree(canvas, tree, x = 0, y = 0)
  return x unless tree

  val_s          = "#{tree.value}"
  val_s_len      = val_s.length
  val_s_len_half = val_s_len / 2

  d_left = x - 2
  d_left = 0 if d_left < 0
  d_left = print_tree(canvas, tree.left,  d_left, y + 2) if tree.left

  sx = d_left + val_s_len + 2
  loop do
    break if put_node?(canvas, val_s, sx, y) && put_node?(canvas, val_s, sx, y - 2)
    sx += 1
  end

  d_right = sx
  d_right  = print_tree(canvas, tree.right,  d_right,  y + 2) if tree.right

  dx = (d_left + d_right) / 2
  sx = dx - val_s_len_half
  loop do
    break if put_node?(canvas, val_s, sx, y) && put_node?(canvas, val_s, sx, y - 2)
    dx += 1
    sx += 1
  end

  if tree.right.nil? && tree.left.nil?
    canvas.put_string(val_s, sx, y)
  else
    connect = val_s
    connect = val_s.center(d_right - d_left - 1, '_') if tree.right && tree.left

    canvas.put_string(connect, d_left + 1,  y)
    canvas.put_string("/",   d_left,  y + 1) if tree.left
    canvas.put_string("\\",  d_right, y + 1) if tree.right
  end
  dx
end

def clone_tree(tree)
  def sub_clone(node)
    return nil if node.nil?
    Node(node.value, sub_clone(node.left), sub_clone(node.right))
  end
  sub_clone(tree)
end

def graph_tree(tree, name = :tree)
  gv = Gviz.new
  gv.graph do
    global layout: :dot
    def sub_graph(node)
      return if node.nil?
      sub_graph(node.left)  unless node.left.nil?
      sub_graph(node.right) unless node.right.nil?

      node node.to_id, label: node.value
      node node.left.to_id, label: node.left.value    unless node.left.nil?
      node node.right.to_id, label: node.right.value  unless node.right.nil?
      if node.left && node.right
        route node.to_id => [node.left.to_id, node.right.to_id]
      elsif !node.left.nil?
        route node.to_id => node.left.to_id
      elsif !node.right.nil?
        route node.to_id => node.right.to_id
      end
    end

    tree2 = clone_tree(tree)
    sub_graph(tree2)
  end
  gv.save(name, :png)
end

canvas = Canvas.new(10, 5)
tree = Node(5,
            Node(4, Node(3, Node(2, Node(1)))),
            Node(9, Node(8, Node(7, Node(5)))))
print_tree(canvas, tree)
canvas.show
puts
graph_tree(tree, :tree_01)

tree = Node(5,
            Node(4, Node(3, Node(2, Node(1)))),
            Node(9, Node(80000,
                         Node(7, Node(5)),
                         Node(3)
                        )
                )
           )
canvas.clear
print_tree(canvas, tree)
canvas.show
puts
graph_tree(tree, :tree_02)

canvas = Canvas.new(10, 6)
tree_s = Node(9,
              Node(4, Node(2), Node(6)),
              Node(15, Node(12), Node(17))
             )
tree = Node(4, tree_s, Node(3, tree_s))
canvas.clear
print_tree(canvas, tree)
canvas.show
puts
graph_tree(tree, :tree_03)

tree = Node(9,
            Node(4, Node(2), Node(6)),
            Node(15, Node(12), Node(17))
           )
canvas.clear
print_tree(canvas, tree)
canvas.show
puts
graph_tree(tree, :tree_04)

__END__
|     5_   |
|    /  \  |
|    4  9  |
|   /  /   |
|   3  8   |
|  /  /    |
|  2  7    |
| /  /     |
| 1  5     |

|     _5__    |
|    /    \   |
|    4    9   |
|   /    /    |
|   3  80000  |
|  /  /     \ |
|  2  7     3 |
| /  /        |
| 1  5        |

|      _____4_____      |
|     /           \     |
|   __9__         3     |
|  /     \       /      |
|  4    15_    __9__    |
| / \  /   \  /     \   |
| 2 6 12  17  4    15_  |
|            / \  /   \ |
|            2 6 12  17 |

|   __9__               |
|  /     \              |
|  4    15_             |
| / \  /   \            |
| 2 6 12  17            |
|                       |
|                       |
|                       |
|                       |

tree_01.png

tree_02.png

tree_03.png

tree_04.png