randomized_binary_search_tree.rb

# ランダム挿入二分探索木
class RandomizedBinarySearchTree
  attr_accessor :value, :left, :right
  attr_accessor :count
  
  # 空の頂点
  def self.empty
    obj = self.new
    obj.value = -1
    obj.count = 0
    obj
  end
  
  @@empty = RandomizedBinarySearchTree.empty

  def initialize(v, l=@@empty, r=@@empty)
    @value = v
    @left = l
    @right = r
    # 自分だけなので部分木のサイズは1
    @count = 1
  end

  # 頂点の追加
  def add(v)
    # 空の場合はそこに追加
    return self.class.new(v)  if self == @@empty
    
    # 1/(部分木サイズ(=count)+1)の確率でこの場所に追加する
    return self.class.add_root(self, v) if rand(@count+1) == 0

    # 後は通常通り
    if v < @value
      @left = @left.add(v)
    else
      @right = @right.add(v)
    end
    
    # 部分木のサイズを更新
    self.update
    
    # 自身を返す
    self
  end
  
  # 要素vを探す
  def search(v)
    return @@empty if self == @@empty
    return self if @value == v
    if v < @value
      # 小さければ左へ
      @left.search(v)
    else
      # 多きければ右へ
      @right.search(v)
    end
  end 

  # 左回転
  def rotate_left
    # 張替え
    tmp = @right
    @right = tmp.left
    tmp.left = self

    # 頂点の部分木サイズを更新する
    self.update
    tmp.update
    
    return tmp
  end
  
  # 右回転
  def rotate_right
    # 張替え
    tmp = @left
    @left = tmp.right
    tmp.right = self
    
    # 頂点の部分木サイズを更新する
    self.update
    tmp.update

    return tmp
  end
    
  # 部分木のサイズ更新。（左の子）+（右の子）+ 1(自分)
  def update
    @count = @left.count + @right.count + 1
  end

  # 木の高さを調べる
  def depth
    return 0 if self == @@empty
    [@left.depth, @right.depth].max + 1
  end
  
  # 根への頂点追加
  def self.add_root(root, v)
    # 最下部にまず追加する
    if root == @@empty
      return self.new(v)
    end

    if v < root.value
      # 自分の左に持ってくる
      root.left = self.add_root(root.left, v)
      
      # 右回転で持ち上げる
      root = root.rotate_right()
    else
      # 自分の右に持ってくる
      root.right = self.add_root(root.right, v)
      
      # 左回転で持ち上げる
      root = root.rotate_left()
    end
    
    # 部分木サイズ更新
    root.update
    
    # 持ち上げた頂点を返す
    root
  end
end

# 1~100を追加
rbst = RandomizedBinarySearchTree.new(1)
(2..100).each do |v|
  rbst = rbst.add(v)
end

# 高さを出力してみる
p rbst.depth