二分探索木の Ruby 実装

binary_search_tree.rb

class Node < Struct.new(:key, :value, :left, :right)
  def inspect
    l, r = "", ""
    l = ", @left='#{left}'"   if left
    r = ", @right='#{right}'" if right
    "#<Node:@key='#{key}', @value='#{value}'#{l}#{r}>"
  end
  alias :to_s :inspect
end

class Tree
  def initialize
    @root = nil
  end
  
  def insert(key, value)
    unless @root
      @root = Node.new(key, value)
      return
    end
    node = @root
    while node
      if key < node.key
        if node.left
          node = node.left
        else
          node.left = Node.new(key, value)
          break
        end
      elsif key > node.key
        if node.right
          node = node.right
        else
          node.right = Node.new(key, value)
          break
        end
      else
        if block_given?
          node.value = yield(key, node.value, value)
        else
          node.value = value
        end
        break
      end
    end
  end
  
  def []=(key, value)
    insert(key, value)
  end
  
  # nodeを返す
  def search(key)
    node = @root
    while node
      if key < node.key
        node = node.left
      elsif key > node.key
        node = node.right
      else
        return node
      end
    end
    nil
  end
  
  # valueを返す
  def [](key)
    search(key)&.value
  end
  
  # Nodeを削除する
  def delete(key)
    delete_min = ->(node) {
      return node.right unless node.left
      node.left = delete_min.(node.left)
      node
    }
    del = ->(node) {
      value = nil
      if node
        if key == node.key
          value = node.value
          if node.left.nil?
            return node.right, value
          elsif node.right.nil?
            return node.left, value
          else
            min_node = search_min(node.right)
            node.key = min_node.key
            node.value = min_node.value
            node.right = delete_min.(node.right)
          end
        elsif key < node.key
          node.left , value = del.(node.left)
        else
          node.right, value = del.(node.right)
        end
      end
      return node, value
    }
    
    @root, value = del.(@root)
    value
  end
  
  # 行きがけ
  def preorder(key = nil, &bk)
    key ||= @root&.key
    return unless key
    traverse = ->(node) {
      if node
        bk.call(node.key, node.value)
        traverse.(node.left)
        traverse.(node.right)
      end
    }
    traverse.(search(key))
  end
  
  # 通りがけ
  def inorder(key = nil, &bk)
    key ||= @root&.key
    return unless key
    traverse = ->(node) {
      if node
        traverse.(node.left)
        bk.call(node.key, node.value)
        traverse.(node.right)
      end
    }
    traverse.(search(key))
  end
  
  # 帰りがけ
  def postorder(key = nil, &bk)
    key ||= @root&.key
    return unless key
    traverse = ->(node) {
      if node
        traverse.(node.left)
        traverse.(node.right)
        bk.call(node.key, node.value)
      end
    }
    traverse.(search(key))
  end
  
  def each(&bk)
    inorder(&bk)
  end
  
  # 幅優先探索
  def breadth_first_search(key = nil, &bk)
    key ||= @root&.key
    return unless key
    queue = [search(key)]
    while (node = queue.shift)
      bk.call(node.key, node.value)
      queue << node.left  if node.left
      queue << node.right if node.right
    end
  end
  alias :bfs :breadth_first_search
  
  def minimum
    if @root
      node = search_min(@root)
      if node
        [node.key, node.value]
      end
    end
  end
  
  def maximum
    search_max = ->(node) {
      node = node.right while node.right
      node
    }
    if @root
      node = search_max.(@root)
      if node
        [node.key, node.value]
      end
    end
  end

  def inspect
    "#<Tree:@root=#{@root}>"
  end
 
  private
  
  def search_min(node)
    node = node.left while node.left
    node
  end
  
  include Enumerable
end

###
### 追加機能
###

class Tree
  def right_rotate(key)
    rotate = ->(node) {
      if node
        if key == node.key
          if node.left
            tmp = node.left
            node.left = tmp.right
            tmp.right = node
            return tmp
          else
            return node
          end
        elsif key < node.key
          node.left  = rotate.(node.left)
        else
          node.right = rotate.(node.right)
        end
      end
      node
    }
    @root = rotate.(@root)
  end
  
  def left_rotate(key)
    rotate = ->(node) {
      if node
        if key == node.key
          if node.right
            tmp = node.right
            node.right = tmp.left
            tmp.left = node
            return tmp
          else
            return node
          end
        elsif key < node.key
          node.left  = rotate.(node.left)
        else
          node.right = rotate.(node.right)
        end
      end
      node
    }
    @root = rotate.(@root)
  end
  
  # nodeを返す
  def parent_by_node(target)
    find = ->(node) {
      if node.left
        return node if target == node.left
        result = find.(node.left)
        return result if result
      end
      if node.right
        return node if target == node.right
        result = find.(node.right)
        return result if result
      end
      nil
    }
    return nil if target == @root
    find.(@root)
  end
  
  # [key, value]を返す
  def parent(key)
    node = search(key)
    node = parent_by_node(node) if node
    node ? [node.key, node.value] : nil
  end
  
  def add_hash(hash)
    hash.each {|k, v| insert(k, v)}
  end
  
  def print
    traverse = ->(node, len = 0) {
      return " --\n" unless node
      st = " -- " + node.key.to_s
      return st + "\n" unless node.left or node.right
      l = st.length
      st += traverse.(node.left, len + l)
      st + " " * (len + l) + traverse.(node.right, len + l)
    }
    puts traverse.(@root)
  end
end