spec/datastructure/sset/red_black_tree_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/red_black_tree"

class SSet::RedBlackTree::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
      (red? && left.red?).should be_false
      left.parent.should eq self
      left.verify
    end

    if right.node?
      key.should be < right.key
      (red? && right.red?).should be_false
      right.parent.should eq self
      right.verify
    end
  end
end

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

run_sset_spec(SSet::RedBlackTree, "SSet::RedBlackTree")

# 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/red_black_tree"
# 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

# Copied with modifications from: https://github.com/crystal-lang/crystal/blob/1.1.1/samples/red_black_tree.cr
class SSet::RedBlackTree(T)
  class Node(T)
    enum Color
      Red
      Black
    end

    include TreeNode(T)

    property key : T, color : Color
    property! left : Node(T), right : Node(T), parent : Node(T)

    def initialize(@key : T, @color : Color = :red)
      @left = @right = @parent = NilNode(T).new
    end

    def black?
      color.black?
    end

    def red?
      color.red?
    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
  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 insert_node(x : Node(T)) : Nil
    insert_helper(x)

    x.color = :red
    while x != root && x.parent.red?
      if x.parent == x.parent.parent.left
        y = x.parent.parent.right
        if !y.nil_node? && y.red?
          x.parent.color = :black
          y.color = :black
          x.parent.parent.color = :red
          x = x.parent.parent
        else
          if x == x.parent.right
            x = x.parent
            left_rotate(x)
          end
          x.parent.color = :black
          x.parent.parent.color = :red
          right_rotate(x.parent.parent)
        end
      else
        y = x.parent.parent.left
        if !y.nil_node? && y.red?
          x.parent.color = :black
          y.color = :black
          x.parent.parent.color = :red
          x = x.parent.parent
        else
          if x == x.parent.left
            x = x.parent
            right_rotate(x)
          end
          x.parent.color = :black
          x.parent.parent.color = :red
          left_rotate(x.parent.parent)
        end
      end
    end
    root.color = :black
  end

  private def delete_node(z : Node(T)) : Nil
    y = (z.left.nil_node? || z.right.nil_node?) ? z : z.succ
    x = y.left.nil_node? ? y.right : y.left
    x.parent = y.parent

    if y.parent.nil_node?
      @root = x
    else
      if y == y.parent.left
        y.parent.left = x
      else
        y.parent.right = x
      end
    end

    z.key = y.key if y != z

    if y.black?
      delete_fixup(x)
    end

    @size -= 1
    y
  end

  def add?(key : T) : Bool
    return false if includes?(key)
    insert_node(Node.new(key))
    true
  end

  def delete(key : T) : Bool
    node = search(key)
    return false if node.nil_node?
    delete_node(node)
    true
  end

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

  def black_height(x : Node = root)
    height = 0
    until x.nil_node?
      x = x.left
      height += 1 if x.nil_node? || x.black?
    end
    height
  end

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

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

  private def left_rotate(x : Node)
    raise "x.right is nil!" if x.right.nil_node?
    y = x.right
    x.right = y.left
    y.left.parent = x if !y.left.nil_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

  private def right_rotate(x : Node)
    raise "x.left is nil!" if x.left.nil_node?
    y = x.left
    x.left = y.right
    y.right.parent = x if !y.right.nil_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

  private def insert_helper(z : Node)
    y = NilNode(T).new
    x = root
    while !x.nil_node?
      y = x
      x = (z.key < x.key) ? x.left : x.right
    end
    z.parent = y
    if y.nil_node?
      @root = z
    else
      z.key < y.key ? y.left = z : y.right = z
    end
    @size += 1
  end

  private def delete_fixup(x : Node)
    while x != root && x.black?
      if x == x.parent.left
        w = x.parent.right
        if w.red?
          w.color = :black
          x.parent.color = :red
          left_rotate(x.parent)
          w = x.parent.right
        end
        if w.left.black? && w.right.black?
          w.color = :red
          x = x.parent
        else
          if w.right.black?
            w.left.color = :black
            w.color = :red
            right_rotate(w)
            w = x.parent.right
          end
          w.color = x.parent.color
          x.parent.color = :black
          w.right.color = :black
          left_rotate(x.parent)
          x = root
        end
      else
        w = x.parent.left
        if w.red?
          w.color = :black
          x.parent.color = :red
          right_rotate(x.parent)
          w = x.parent.left
        end
        if w.right.black? && w.left.black?
          w.color = :red
          x = x.parent
        else
          if w.left.black?
            w.right.color = :black
            w.color = :red
            left_rotate(w)
            w = x.parent.left
          end
          w.color = x.parent.color
          x.parent.color = :black
          w.left.color = :black
          right_rotate(x.parent)
          x = root
        end
      end
    end
    x.color = :black
  end
end

class SSet::RedBlackTree::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
      (red? && left.red?).should be_false
      left.parent.should eq self
      left.verify
    end

    if right.node?
      key.should be < right.key
      (red? && right.red?).should be_false
      right.parent.should eq self
      right.verify
    end
  end
end

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

run_sset_spec(SSet::RedBlackTree, "SSet::RedBlackTree")