破壊操作型 AVL 木

制約条件「すべてのノードについて、左右の部分木の高さの差を 1 に収めた」木が AVL 木です。1962 年に発表された AVL 木は、挿入・削除・探索のすべてで O(log n) を満たす最初に発見された 2 分探索木であり、それらを満たす最も制約が厳しい 2 分探索木です。歴史的価値があり、バランス調整がエレガントに記述できるとあって、教科書向けの美しさを誇ります。ですが、世の常として、美しさは融通性の犠牲によって得たものであり、泥臭い赤黒木よりも使い道を選り好みします。

一本の木への挿入・削除・探索しかおこなわない、そんな気高い応用に AVL 木は向いています。その手の破壊操作型 AVL 木をループで実装している優れた実例が、4.3BSD の csh の変数辞書です。

⇒ http://minnie.tuhs.org/cgi-bin/utree.pl?file=4.3BSD/usr/src/bin/csh/sh.set.c

この実装は、挿入と削除の両方に同じバランス調整ルーチンを使っている上に古い C 言語の悪い癖も匠に利用してあって、記述のエレガントにほれぼれできます。今回は、このバランス調整ルーチンを元に AVL 木を書いてみることにします。ただし、実用上は問題ないものの、削除により左右が NULL になったノードのバランス・マークがアンバランスのままで放置される挙動があり、デバッグするときに私のような凡人は惑わされたこともあったので、そこは修正しています。なお、回転はノードを入れ替える本来のやりかたではなく、ノードの値とリンクをごっそり入れ替える別のやりかたに変えています。

class AVLtree
  class Node
    attr_accessor :bal, :left, :data, :right

    def initialize(bal, left, data, right)
      @bal, @left, @data, @right = bal, left, data, right
    end

    def inspect
      '(%s %s %s)' % [@left.inspect, @data.inspect, @right.inspect]
    end
  end

  def inspect() @root.inspect end

  attr_accessor :root

  def initialize(root = nil)
    @root = root
  end

  def insert(x)
    child = Node.new(0, nil, x, nil)
    return self.tap{ @root = child } if @root.nil?
    path = []
    node = @root
    while not node.nil?
      return self if x == node.data
      path.push node
      if x < node.data
        node = node.left
      else
        node = node.right
      end
    end
    path.push child
    parent = path[-2]
    if x < parent.data
      parent.left = child
    else
      parent.right = child
    end
    balance(path, :insert)
    self
  end

  def delete(x)
    return self if @root.nil?
    path = []
    just, node = nil, @root
    while not node.nil?
      if x == node.data
        just = node
      end
      path.push node
      if x <= node.data
        node = node.left
      else
        node = node.right
      end
    end
    return self if just.nil?
    node = path[-1]     # path = [..., parent, node]
    parent = path[-2]
    if not just.equal?(node)
      just.data = node.data
    end
    child = if node.right.nil? then node.left else node.right end
    if parent.nil?
      @root = child
    elsif parent.left.equal?(node)
      parent.left = child
    else
      parent.right = child
    end
    path[-1] = child
    balance(path, :delete)
    self
  end

  # based on csh/sh.set.c 5.2 (Berkeley) 1985-06-06
  # see http://minnie.tuhs.org/cgi-bin/utree.pl?file=4.3BSD/usr/src/bin/csh/sh.set.c

  def balance(path, d)
    path.each_cons(2).reverse_each do |n, t|
      if (d == :delete) ^ (n.right.equal?(t))
        balance_right_heavy(n)
      else
        balance_left_heavy(n)
      end
      break if (d == :delete) ^ (n.bal == 0)
    end
  end

  def balance_right_heavy(n)
    if n.bal == 0
      n.bal = +1
    elsif n.bal < 0 or n.right.nil?  # was left heavy
      n.bal =  0
    else  # was right heavy
      case n.right.bal
      when +1
        rotate_left(n)
        n.left.bal = 0
        n.bal = 0
      when  0
        rotate_left(n)
        n.left.bal = +1
        n.bal = -1
      when -1
        b = n.right.left.bal
        rotate_right(n.right)
        rotate_left(n)
        n.left.bal  = if b > 0 then -1 else 0 end
        n.right.bal = if b < 0 then +1 else 0 end
        n.bal = 0
      end
    end
  end

  def balance_left_heavy(n)
    if n.bal == 0
      n.bal = -1
    elsif n.bal > 0 or n.left.nil?  # was right heavy
      n.bal =  0
    else  # was left heavy
      case n.left.bal
      when -1
        rotate_right(n)
        n.right.bal = 0
        n.bal = 0
      when  0
        rotate_right(n)
        n.right.bal = -1
        n.bal = +1
      when +1
        b = n.left.right.bal
        rotate_left(n.left)
        rotate_right(n)
        n.left.bal  = if b > 0 then -1 else 0 end
        n.right.bal = if b < 0 then +1 else 0 end
        n.bal = 0
      end
    end
  end

  def rotate_left(n)
    # #n=(p0 d1 #m=(p2 d3 p4)) => #n=(#m=(p0 d1 p2) d3 p4)
    p0, d1, p2, d3, p4 = n.left, n.data, n.right.left, n.right.data, n.right.right
    n.left = n.right
    n.left.left, n.left.data, n.left.right, n.data, n.right = p0, d1, p2, d3, p4
    n
  end

  def rotate_right(n)
    # #n=(#m=(p0 d1 p2) d3 p4) => #n=(p0 d1 #m=(p2 d3 p4))
    p0, d1, p2, d3, p4 = n.left.left, n.left.data, n.left.right, n.data, n.right
    n.right = n.left
    n.left, n.data, n.right.left, n.right.data, n.right.right = p0, d1, p2, d3, p4
  end
end

tree = AVLtree.new
[5, 3, 8, 2, 4, 7, 10, 1, 6, 9, 11].each do |x|
  p tree.insert(x)
end
[4, 8, 6, 5, 2, 1, 7].each do |x|
  p tree.delete(x)
end