spec/datastructure/sset/treap_spec.cr

View this file on GitHub
Last update: 2022-01-08 20:48:10+09:00

Depends on

Code

require "./spec_helper"
require "../../../src/datastructure/sset/treap"

class SSet::Treap::Node
  def verify
    if nil_node?
      left.nil_node?.should be_true
      right.nil_node?.should be_true
    end

    if left.node?
      key.should be > left.key
      priority.should be <= left.priority
      left.parent.should eq self
      left.verify
    end

    if right.node?
      key.should be < right.key
      priority.should be <= right.priority
      right.parent.should eq self
      right.verify
    end
  end
end

class SSet::Treap
  def verify
    if root.node?
      root.parent.nil_node?.should be_true
    end
    root.verify
  end
end

run_sset_spec(SSet::Treap, "SSet::Treap")

# require "./spec_helper"
require "spec"

def run_sset_spec(type : SSet.class, class_name : String) forall SSet
  describe class_name + "(Int32)" do
    it "{}" do
      SSet{3, 1, 4, 1, 5}.to_a.should eq [1, 3, 4, 5]
    end

    it "#root" do
      s = SSet(Int32).new
      s.root.nil_node?.should be_true
    end

    it "#size" do
      s = SSet(Int32).new
      s.size.should eq 0
      s.add 1
      s.size.should eq 1
      s.add 1
      s.size.should eq 1
      s.delete 1
      s.size.should eq 0
    end

    it "#empty?" do
      s = SSet(Int32).new
      s.empty?.should be_true
      s.add 1
      s.empty?.should be_false
    end

    it "#clear" do
      s = SSet(Int32).new
      s.add 1
      s.clear.size.should eq 0
    end

    it "#add?" do
      s = SSet(Int32).new
      s.add?(1).should be_true
      s.add?(1).should be_false
      s.verify
    end

    it "#add, #<<" do
      s = SSet(Int32).new
      s.add(1).add(2).add(1)
      s.size.should eq 2
      s << 3 << 4 << 3
      s.size.should eq 4
      s.verify
    end

    it "#min_node, #max_node" do
      s = SSet(Int32).new
      s.min_node.nil_node?.should be_true
      s.max_node.nil_node?.should be_true
      s << 1 << 2
      s.min_node.key?.should eq 1
      s.max_node.key?.should eq 2
    end

    it "#min?, #min, #max?, #max" do
      s = SSet(Int32).new
      s.min?.should be_nil
      s.max?.should be_nil
      expect_raises(SSet::EmptyError) { s.min }
      expect_raises(SSet::EmptyError) { s.max }
      s << 1 << 2
      s.min?.should eq 1
      s.max?.should eq 2
      s.min.should eq 1
      s.max.should eq 2
    end

    it "#split" do
      1000.times do
        s = SSet.new(1..10)
        l, r = s.split(5)
        l.to_a.should eq [1, 2, 3, 4, 5]
        r.to_a.should eq [6, 7, 8, 9, 10]
        s.verify
      end
      10.times do
        s = SSet.new(1..1000)
        l, r = s.split(500)
        l.to_a.should eq (1..500).to_a
        r.to_a.should eq (501..1000).to_a
        s.verify
      end
    end

    it "#each" do
      s = SSet{3, 1, 4, 1, 5}

      a = [] of Int32
      s.each { |x| a << x }
      a.should eq [1, 3, 4, 5]

      b = [] of Int32
      s.each(s.search(3)) { |x| b << x }
      b.should eq [3, 4, 5]

      s.each.to_a.should eq [1, 3, 4, 5]
      s.each(s.search(3)).to_a.should eq [3, 4, 5]
    end

    it "#reverse_each" do
      s = SSet{3, 1, 4, 1, 5}

      a = [] of Int32
      s.reverse_each { |x| a << x }
      a.should eq [5, 4, 3, 1]

      a = [] of Int32
      s.reverse_each(s.search(4)) { |x| a << x }
      a.should eq [4, 3, 1]

      s.reverse_each.to_a.should eq [5, 4, 3, 1]
      s.reverse_each(s.search(4)).to_a.should eq [4, 3, 1]
    end

    it "#includes?" do
      s = SSet{1, 3, 5}
      {1, 3, 5}.each { |x| s.includes?(x).should be_true }
      {0, 2, 4}.each { |x| s.includes?(x).should be_false }
    end

    it "#search" do
      s = SSet{1, 3, 5}
      {1, 3, 5}.each { |x| s.search(x).key?.should eq x }
      {0, 2, 4}.each { |x| s.search(x).nil_node?.should be_true }
    end

    it "#le, #le!" do
      s = SSet{1, 3}
      [nil, 1, 1, 3, 3].each_with_index { |e, x| s.le(x).should eq e }
      expect_raises(NilAssertionError) { s.le!(0) }
      [1, 1, 3, 3].each_with_index(1) { |e, x| s.le!(x).should eq e }
    end

    it "#lt, #lt!" do
      s = SSet{1, 3}
      [nil, nil, 1, 1, 3, 3].each_with_index { |e, x| s.lt(x).should eq e }
      expect_raises(NilAssertionError) { s.lt!(0) }
      expect_raises(NilAssertionError) { s.lt!(1) }
      [1, 1, 3, 3].each_with_index(2) { |e, x| s.lt!(x).should eq e }
    end

    it "#ge, #ge!" do
      s = SSet{1, 3}
      [1, 1, 3, 3, nil].each_with_index { |e, x| s.ge(x).should eq e }
      expect_raises(NilAssertionError) { s.ge!(4) }
      [1, 1, 3, 3].each_with_index { |e, x| s.ge!(x).should eq e }
    end

    it "#gt, #gt!" do
      s = SSet{1, 3}
      [1, 3, 3, nil, nil].each_with_index { |e, x| s.gt(x).should eq e }
      expect_raises(NilAssertionError) { s.gt!(3) }
      expect_raises(NilAssertionError) { s.gt!(4) }
      [1, 3, 3].each_with_index { |e, x| s.gt!(x).should eq e }
    end

    it "include Enumerable" do
      s = SSet{1, 2, 3, 4}
      s.all? { |x| x > 0 }.should be_true
    end

    it "include Iterable" do
      s = SSet{1, 2, 3, 4}
      s.cycle(2).to_a.should eq [1, 2, 3, 4, 1, 2, 3, 4]
    end

    it "#to_s, #inspect" do
      s = SSet{1, 2, 3, 4}
      s.to_s.should eq "#{class_name}{1, 2, 3, 4}"
      s.inspect.should eq "#{class_name}{1, 2, 3, 4}"
    end

    it "big" do
      n = 10**4
      s = SSet(Int32).new
      (1..n).each do |x|
        s << x
        s.size.should eq x
        s.min.should eq 1
        s.max.should eq x
        s.verify
      end
      s.to_a.should eq (1..n).to_a
      (-n..n*2).each do |x|
        s.includes?(x).should(1 <= x <= n ? be_true : be_false)
      end
      (1..n).each do |x|
        s.delete(x).should be_true
        s.delete(x).should be_false
      end
    end
  end

  it class_name + "(String)" do
    SSet{"a", "c", "b"}.to_a.should eq %w[a b c]
    SSet{"a", "ab", "abc", "abc"}.to_a.should eq %w[a ab abc]
  end
end

# require "../../../src/datastructure/sset/treap"
# require "../binary_tree"
module TreeNode(T)
  abstract def key : T
  abstract def left
  abstract def right
  abstract def parent

  def node?
    true
  end

  def nil_node?
    false
  end

  def key? : T?
    key
  end

  def key! : T
    key
  end

  def min_node : self
    x = self
    while x.left.node?
      x = x.left
    end
    x
  end

  def max_node : self
    x = self
    while x.right.node?
      x = x.right
    end
    x
  end

  def succ : self
    if right.node?
      return right.min_node
    end
    x, y = self, parent
    while y.node? && x != y.left
      x = y
      y = y.parent
    end
    y
  end

  def pred : self
    if left.node?
      return left.max_node
    end
    x, y = self, parent
    while y.node? && x != y.right
      x = y
      y = y.parent
    end
    y
  end

  def to_s(io : IO)
    io << '[' << key << ']'
  end

  def inspect(io : IO)
    to_s(io)
  end
end

module TreeNilNode(T)
  def node?
    false
  end

  def nil_node?
    true
  end

  def key? : T?
    nil
  end

  def key!
    raise NilAssertionError.new
  end

  def to_s(io : IO)
    io << "NilNode"
  end

  def inspect(io : IO)
    to_s(io)
  end
end

module BinaryTree(T, Node, NilNode)
  class EmptyError < Exception
    def initialize(message = "Empty RedBlackTree")
      super(message)
    end
  end

  include Enumerable(T)
  include Iterable(T)

  abstract def root
  abstract def size : Int32
  abstract def add?(key : T) : Bool
  abstract def delete(key : T) : Bool

  def empty? : Bool
    root.nil_node?
  end

  def clear : self
    @root = NilNode.new
    @size = 0
    self
  end

  def add(key : T) : self
    add?(key)
    self
  end

  def <<(key : T) : self
    add(key)
  end

  def rotate_left(x : Node) : Nil
    raise "x.right is nil!" if x.right.nil_node?
    y = x.right
    x.right = y.left
    y.left.parent = x if y.left.node?
    y.parent = x.parent
    if x.parent.nil_node?
      @root = y
    else
      if x == x.parent.left
        x.parent.left = y
      else
        x.parent.right = y
      end
    end
    y.left = x
    x.parent = y
  end

  def rotate_right(x : Node) : Nil
    raise "x.left is nil!" if x.left.nil_node?
    y = x.left
    x.left = y.right
    y.right.parent = x if y.right.node?
    y.parent = x.parent
    if x.parent.nil_node?
      @root = y
    else
      if x == x.parent.left
        x.parent.left = y
      else
        x.parent.right = y
      end
    end
    y.right = x
    x.parent = y
  end

  def min_node : Node
    root.min_node
  end

  def max_node : Node
    root.max_node
  end

  def min? : T?
    min_node.key?
  end

  def min : T
    min? || raise EmptyError.new
  end

  def max? : T?
    max_node.key?
  end

  def max : T
    max? || raise EmptyError.new
  end

  private class InorderWalkIterator(T, Node)
    include Iterator(T)

    def initialize(@node : Node)
    end

    def next
      if @node.nil_node?
        stop
      else
        @node.key.tap { @node = @node.succ }
      end
    end
  end

  def each(node : Node = min_node, &block) : Nil
    while node.node?
      yield node.key
      node = node.succ
    end
  end

  def each(node : Node = min_node)
    InorderWalkIterator(T, Node).new(node)
  end

  private class ReverseInorderWalkIterator(T, Node)
    include Iterator(T)

    def initialize(@node : Node)
    end

    def next
      if @node.nil_node?
        stop
      else
        @node.key.tap { @node = @node.pred }
      end
    end
  end

  def reverse_each(node : Node = max_node, &block) : Nil
    while node.node?
      yield node.key
      node = node.pred
    end
  end

  def reverse_each(node : Node = max_node)
    ReverseInorderWalkIterator(T, Node).new(node)
  end

  def includes?(key : T) : Bool
    x = root
    loop do
      return false if x.nil_node?
      cmp = key <=> x.key
      return true if cmp == 0
      x = cmp < 0 ? x.left : x.right
    end
  end

  def search(key : T, x : Node = root) : Node
    while x.node?
      cmp = key <=> x.key
      break if cmp == 0
      x = cmp < 0 ? x.left : x.right
    end
    x
  end

  def le_node(key : T, x : Node = root) : Node
    loop do
      cmp = key <=> x.key
      y = cmp < 0 ? x.left : x.right
      if y.nil_node?
        return cmp < 0 ? x.pred : x
      end
      x = y
    end
  end

  def lt_node(key : T, x : Node = root) : Node
    loop do
      cmp = key <=> x.key
      y = cmp <= 0 ? x.left : x.right
      if y.nil_node?
        return cmp <= 0 ? x.pred : x
      end
      x = y
    end
  end

  def ge_node(key : T, x : Node = root) : Node
    loop do
      cmp = key <=> x.key
      y = cmp <= 0 ? x.left : x.right
      if y.nil_node?
        return cmp <= 0 ? x : x.succ
      end
      x = y
    end
  end

  def gt_node(key : T, x : Node = root) : Node
    loop do
      cmp = key <=> x.key
      y = cmp < 0 ? x.left : x.right
      if y.nil_node?
        return cmp < 0 ? x : x.succ
      end
      x = y
    end
  end

  {% for method in [:le, :lt, :ge, :gt] %}
    def {{method.id}}(key : T) : T?
      {{method.id}}_node(key).key?
    end

    def {{method.id}}!(key : T) : T
      {{method.id}}_node(key).key!
    end
  {% end %}
end

class SSet::Treap(T)
  class Node(T)
    include TreeNode(T)

    getter key : T, priority : Int32
    property! left : Node(T), right : Node(T), parent : Node(T)

    def initialize(@key : T, @priority : Int32)
      @left = @right = @parent = NilNode(T).new
    end

    def split(split_key : T) : {Node(T), Node(T)}
      if nil_node?
        {NilNode(T).new, NilNode(T).new}
      elsif split_key < key
        l, r = left.split(split_key)
        self.left = r
        r.parent = self
        {l, self}
      else
        l, r = right.split(split_key)
        self.right = l
        l.parent = self
        {self, r}
      end
    end

    def to_s(io : IO)
      io << '[' << key << ']'
    end
  end

  class NilNode(T) < Node(T)
    include TreeNilNode(T)

    def initialize
      @key = uninitialized T
      @left = @right = @parent = self
    end
  end

  include BinaryTree(T, Node(T), NilNode(T))
  getter root : Node(T), size : Int32 = 0

  def initialize
    @root = NilNode(T).new
  end

  def initialize(enumerable : Enumerable(T))
    initialize
    enumerable.each { |x| self << x }
  end

  protected def initialize(@root : Node(T))
  end

  private def add_node(node : Node) : Bool
    if empty?
      @root = node
      @size += 1
      return true
    end
    u, prev = root, NilNode(T).new
    while u.node?
      cmp = node.key <=> u.key
      return false if cmp == 0
      u, prev = cmp < 0 ? u.left : u.right, u
    end
    if node.key < prev.key
      prev.left = node
    else
      prev.right = node
    end
    node.parent = prev
    @size += 1
    true
  end

  private def bubble_up(u : Node) : Nil
    raise "u is nil node" unless u.node?
    while u.parent.node? && u.parent.priority > u.priority
      if u.parent.right == u
        rotate_left(u.parent)
      else
        rotate_right(u.parent)
      end
    end
    @root = u if u.parent.nil_node?
  end

  def add?(key : T) : Bool
    u = Node.new(key, rand(Int32))
    add_node(u).tap do |f|
      bubble_up(u) if f
    end
  end

  private def trickle_down(u : Node(T))
    loop do
      l, r = u.left.node?, u.right.node?
      break unless l || r
      if !l
        rotate_left(u)
      elsif !r
        rotate_right(u)
      elsif u.left.priority < u.right.priority
        rotate_right(u)
      else
        rotate_left(u)
      end
      @root = u.parent if root == u
    end
  end

  def delete(key : T) : Bool
    u = search(key)
    return false if u.nil_node?
    trickle_down(u)
    if u.parent.left == u
      u.parent.left = NilNode(T).new
    else
      u.parent.right = NilNode(T).new
    end
    @size -= 1
    if size == 0
      @root = NilNode(T).new
    end
    true
  end

  def split(key : T) : {self, self}
    l, r = root.split(key)
    l.parent = NilNode(T).new
    r.parent = NilNode(T).new
    {Treap(T).new(l), Treap(T).new(r)}
  end

  def to_s(io : IO) : Nil
    io << "SSet::Treap{"
    each_with_index do |x, i|
      io << ", " if i > 0
      io << x
    end
    io << '}'
  end

  def inspect(io : IO) : Nil
    to_s(io)
  end
end

class SSet::Treap::Node
  def verify
    if nil_node?
      left.nil_node?.should be_true
      right.nil_node?.should be_true
    end

    if left.node?
      key.should be > left.key
      priority.should be <= left.priority
      left.parent.should eq self
      left.verify
    end

    if right.node?
      key.should be < right.key
      priority.should be <= right.priority
      right.parent.should eq self
      right.verify
    end
  end
end

class SSet::Treap
  def verify
    if root.node?
      root.parent.nil_node?.should be_true
    end
    root.verify
  end
end

run_sset_spec(SSet::Treap, "SSet::Treap")