Treap with order statistics

treap.rb

class Treap
  class Node
    attr_accessor :key, :priority, :left, :right, :count
    
    def initialize(key, priority)
      @key = key
      @priority = priority
      @count = 1
    end
    
    def self.count_of(node)
      return node ? node.count : 0
    end
    
    def update_count
      @count = Node.count_of(@left) + Node.count_of(@right) + 1
    end
    
    def kth(k)
      if k == Node.count_of(@left)
        return @key
      elsif k < Node.count_of(@left)
        return @left.kth(k)
      else
        return @right.kth(k - Node.count_of(@left) - 1)
      end
    end
    
    def self.merge(node1, node2)
      if not node1 or not node2
        return node1 ? node1 : node2
      elsif node1.priority > node2.priority
        node1.right = merge(node1.right, node2)
        node1.update_count
        return node1
      else
        node2.left = merge(node1, node2.left)
        node2.update_count
        return node2
      end
    end
    
    def self.split(node, key)
      if not node
        return [nil, nil]
      elsif node.key <= key
        spliting_right_result = split(node.right, key)
        node.right = spliting_right_result[0]
        node.update_count
        return [node, spliting_right_result[1]]
      else
        spliting_left_result = split(node.left, key)
        node.left = spliting_left_result[1]
        node.update_count
        return [spliting_left_result[0], node]
      end
    end
    
    def self.traverse(node)
      if not node
        return []
      end
      result = []
      if node.left
        result += traverse(node.left)
      end
      result += [node.key]
      if node.right
        result += traverse(node.right)
      end
      return result
    end
    
    def self.check(node, key)
      if not node
        return false
      elsif node.key == key
        return true
      elsif node.key < key
        return check(node.left, key)
      elsif node.key > key
        return check(node.right, key)
      end
    end
  end
  
  attr_accessor :root
  
  def initialize
    @root = nil
  end
  
  def insert(key)
    left, right = Node.split(@root, key)
    new_node = Node.new(key, rand(2000000000))
    @root = Node.merge(left, Node.merge(new_node, right))
  end
  
  def remove(key)
    left_result_part, rest = Node.split(@root, key - 1)
    elements_to_remove, right_result_part = Node.split(rest, key)
    @root = Node.merge(left_result_part, right_result_part)
  end
  
  def traverse
    return Node.traverse(@root)
  end
  
  def check(key)
    return Node.check(@root, key)
  end
  
  def kth(k)
    return @root.kth(k)
  end
end