This documentation is automatically generated by online-judge-tools/verification-helper

:warning: src/datastructure/smultiset/red_black_tree.cr

Depends on

Required by

Code

require "../binary_tree"

# Copied with modifications from: https://github.com/crystal-lang/crystal/blob/1.1.1/samples/red_black_tree.cr
class SMultiSet::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
    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 << "SMultiSet::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
# 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 SMultiSet::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
    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 << "SMultiSet::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
Back to top page