Red black tree sort/Julia

From Rosetta Code

See Red_black_tree_sort

Code originally by Koustav Chowdhury.

""" Modified from JuliaCollections/DataStructures.jl's file red_black_tree.jl """

# leftChild has keys which are less than the node
# rightChild has keys which are greater than the node
# isred is true if it's a Red Node, else it's false
# keys are unique
mutable struct RBTreeNode{K}
    isred::Bool
    data::Union{K, Nothing}
    leftChild::Union{Nothing, RBTreeNode{K}}
    rightChild::Union{Nothing, RBTreeNode{K}}
    parent::Union{Nothing, RBTreeNode{K}}

    RBTreeNode{K}() where K = new{K}(true, nothing, nothing, nothing, nothing)

    RBTreeNode{K}(d::K) where K = new{K}(true, d, nothing, nothing, nothing)
end

RBTreeNode() = RBTreeNode{Any}()
RBTreeNode(d) = RBTreeNode{Any}(d)

function create_null_node(K::Type)
    node = RBTreeNode{K}()
    node.isred = false
    return node
end

mutable struct RBTree{K}
    root::RBTreeNode{K}
    nil::RBTreeNode{K}
    count::Int

    function RBTree{K}() where K
        rb = new()
        rb.nil = create_null_node(K)
        rb.root = rb.nil
        rb.count = 0
        return rb
    end
end

RBTree() = RBTree{Any}()

Base.length(tree::RBTree) = tree.count

"""
    search_node(tree, key)

Returns the last visited node, while traversing through in binary-search-tree fashion looking for `key`.
"""
function search_node(tree::RBTree{K}, d::K) where K
    node = tree.root
    while node !== tree.nil && d != node.data
        if d < node.data
            node = node.leftChild
        else
            node = node.rightChild
        end
    end
    return node
end

"""
    haskey(tree, key)

Returns true if `key` is present in the `tree`, else returns false.
"""
function Base.haskey(tree::RBTree{K}, d::K) where K
    node = search_node(tree, d)
    return (node.data == d)
end

"""
    insert_node!(tree::RBTree, node::RBTreeNode)

Inserts `node` at proper location by traversing through the `tree` in a binary-search-tree fashion.
"""
function insert_node!(tree::RBTree, node::RBTreeNode)
    node_y = nothing
    node_x = tree.root

    while node_x !== tree.nil
        node_y = node_x
        if node.data < node_x.data
            node_x = node_x.leftChild
        else
            node_x = node_x.rightChild
        end
    end

    node.parent = node_y
    if node_y == nothing
        tree.root = node
    elseif node.data < node_y.data
        node_y.leftChild = node
    else
        node_y.rightChild = node
    end
end

"""
    left_rotate!(tree::RBTree, node_x::RBTreeNode)

Performs a left-rotation on `node_x` and updates `tree.root`, if required.
"""
function left_rotate!(tree::RBTree, node_x::RBTreeNode)
    node_y = node_x.rightChild
    node_x.rightChild = node_y.leftChild
    if node_y.leftChild !== tree.nil
        node_y.leftChild.parent = node_x
    end
    node_y.parent = node_x.parent
    if (node_x.parent == nothing)
        tree.root = node_y
    elseif (node_x == node_x.parent.leftChild)
        node_x.parent.leftChild = node_y
    else
        node_x.parent.rightChild = node_y
    end
    node_y.leftChild = node_x
    node_x.parent = node_y
end

"""
    right_rotate!(tree::RBTree, node_x::RBTreeNode)

Performs a right-rotation on `node_x` and updates `tree.root`, if required.
"""
function right_rotate!(tree::RBTree, node_x::RBTreeNode)
    node_y = node_x.leftChild
    node_x.leftChild = node_y.rightChild
    if node_y.rightChild !== tree.nil
        node_y.rightChild.parent = node_x
    end
    node_y.parent = node_x.parent
    if (node_x.parent == nothing)
        tree.root = node_y
    elseif (node_x == node_x.parent.leftChild)
        node_x.parent.leftChild = node_y
    else
        node_x.parent.rightChild = node_y
    end
    node_y.rightChild = node_x
    node_x.parent = node_y
end

"""
   fix_insert!(tree::RBTree, node::RBTreeNode)

This method is called to fix the property of having no two adjacent nodes of red color in the `tree`.
"""
function fix_insert!(tree::RBTree, node::RBTreeNode)
    parent = nothing
    grand_parent = nothing
    # for root node, we need to change the color to black
    # other nodes, we need to maintain the property such that
    # no two adjacent nodes are red in color
    while  node != tree.root && node.parent.isred
        parent = node.parent
        grand_parent = parent.parent

        if (parent == grand_parent.leftChild) # parent is the leftChild of grand_parent
            uncle = grand_parent.rightChild

            if (uncle.isred) # uncle is red in color
                grand_parent.isred = true
                parent.isred = false
                uncle.isred = false
                node = grand_parent
            else  # uncle is black in color
                if (node == parent.rightChild) # node is rightChild of its parent
                    node = parent
                    left_rotate!(tree, node)
                end
                # node is leftChild of its parent
                node.parent.isred = false
                node.parent.parent.isred = true
                right_rotate!(tree, node.parent.parent)
            end
        else # parent is the rightChild of grand_parent
            uncle = grand_parent.leftChild

            if (uncle.isred) # uncle is red in color
                grand_parent.isred = true
                parent.isred = false
                uncle.isred = false
                node = grand_parent
            else  # uncle is black in color
                if (node == parent.leftChild) # node is leftChild of its parent
                    node = parent
                    right_rotate!(tree, node)
                end
                # node is rightChild of its parent
                node.parent.isred = false
                node.parent.parent.isred = true
                left_rotate!(tree, node.parent.parent)
            end
        end
    end
    tree.root.isred = false
end

"""
    insert!(tree, key)

Inserts `key` in the `tree` if it is not present.
"""
function Base.insert!(tree::RBTree{K}, d::K) where K
    # if the key exists in the tree, no need to insert
    haskey(tree, d) && return tree

    # insert, if not present in the tree
    node = RBTreeNode{K}(d)
    node.leftChild = node.rightChild = tree.nil

    insert_node!(tree, node)

    if node.parent == nothing
        node.isred = false
    elseif node.parent.parent == nothing
        ;
    else
        fix_insert!(tree, node)
    end
    tree.count += 1
    return tree
end

"""
    delete_fix(tree::RBTree, node::Union{RBTreeNode, Nothing})

This method is called when a black node is deleted because it violates the black depth property of the RBTree.
"""
function delete_fix(tree::RBTree, node::Union{RBTreeNode, Nothing})
    while node != tree.root && !node.isred
         if node == node.parent.leftChild
            sibling = node.parent.rightChild
            if sibling.isred
                sibling.isred = false
                node.parent.isred = true
                left_rotate!(tree, node.parent)
                sibling = node.parent.rightChild
            end

            if !sibling.rightChild.isred && !sibling.leftChild.isred
                sibling.isred = true
                node = node.parent
            else
                if !sibling.rightChild.isred
                    sibling.leftChild.isred = false
                    sibling.isred = true
                    right_rotate!(tree, sibling)
                    sibling = node.parent.rightChild
                end

                sibling.isred = node.parent.isred
                node.parent.isred = false
                sibling.rightChild.isred = false
                left_rotate!(tree, node.parent)
                node = tree.root
            end
        else
            sibling = node.parent.leftChild
            if sibling.isred
                sibling.isred = false
                node.parent.isred = true
                right_rotate!(tree, node.parent)
                sibling = node.parent.leftChild
            end

            if !sibling.rightChild.isred && !sibling.leftChild.isred
                sibling.isred = true
                node = node.parent
            else
                if !sibling.leftChild.isred
                    sibling.rightChild.isred = false
                    sibling.isred = true
                    left_rotate!(tree, sibling)
                    sibling = node.parent.leftChild
                end

                sibling.isred = node.parent.isred
                node.parent.isred = false
                sibling.leftChild.isred = false
                right_rotate!(tree, node.parent)
                node = tree.root
            end
        end
    end
    node.isred = false
    return nothing
end

"""
    rb_transplant(tree::RBTree, u::Union{RBTreeNode, Nothing}, v::Union{RBTreeNode, Nothing})

Replaces `u` by `v` in the `tree` and updates the `tree` accordingly.
"""
function rb_transplant(tree::RBTree, u::Union{RBTreeNode, Nothing}, v::Union{RBTreeNode, Nothing})
    if u.parent == nothing
        tree.root = v
    elseif u == u.parent.leftChild
        u.parent.leftChild = v
    else
        u.parent.rightChild = v
    end
    v.parent = u.parent
end

"""
   minimum_node(tree::RBTree, node::RBTreeNode)

Returns the RBTreeNode with minimum value in subtree of `node`.
"""
function minimum_node(tree::RBTree, node::RBTreeNode)
    (node === tree.nil) && return node
    while node.leftChild !== tree.nil
        node = node.leftChild
    end
    return node
end

"""
    delete!(tree::RBTree, key)

Deletes `key` from `tree`, if present, else returns the unmodified tree.
"""
function Base.delete!(tree::RBTree{K}, d::K) where K
    z = tree.nil
    node = tree.root

    while node !== tree.nil
        if node.data == d
            z = node
        end

        if d < node.data
            node = node.leftChild
        else
            node = node.rightChild
        end
    end

    (z === tree.nil) && return tree

    y = z
    y_original_color = y.isred
    x = RBTreeNode{K}()
    if z.leftChild === tree.nil
        x = z.rightChild
        rb_transplant(tree, z, z.rightChild)
    elseif z.rightChild === tree.nil
        x = z.leftChild
        rb_transplant(tree, z, z.leftChild)
    else
        y = minimum_node(tree, z.rightChild)
        y_original_color = y.isred
        x = y.rightChild

        if y.parent == z
            x.parent = y
        else
            rb_transplant(tree, y, y.rightChild)
            y.rightChild = z.rightChild
            y.rightChild.parent = y
        end

        rb_transplant(tree, z, y)
        y.leftChild = z.leftChild
        y.leftChild.parent = y
        y.isred = z.isred
    end

    !y_original_color && delete_fix(tree, x)
    tree.count -= 1
    return tree
end

"""
    function sortednodes(node::RBTreeNode{K}) where K

Return a sorted node list. Recursive.
"""
function sortednodes(node::RBTreeNode{K}, nilnode) where K
    if (node !== nilnode)
        left = sortednodes(node.leftChild, nilnode)
        right = sortednodes(node.rightChild, nilnode)
        return append!(push!(left, node), right)
    else
        return RBTreeNode{K}[]
    end
end

"""
   Print the RBTree (red black tree) as a list of nodes with their keys.
"""
function Base.print(io::IO, tree::RBTree)
    println(io, "Red-Black tree node list is:")
    for node in sortednodes(tree.root, tree.nil)
        left = node.leftChild === tree.nil ? "empty" : string(node.leftChild.data)
        color = node.isred ? "red" : "black"
        right = node.rightChild === tree.nil ? "empty" : string(node.rightChild.data)
        parent = node.parent isa Nothing ? "empty" : string(node.parent.data)
        println(io, "(left: $left, color: $color, current: $(node.data), right: $right, parent: $parent)")
    end
end

function RBTsort()
    tree, data = RBTree(), Int[]
    while length(unique(data)) < 30
        data = rand(1:1000, 30)
    end
    for i in data
        insert!(tree, i)
    end
    println("After adding 30 nodes, ", tree)
    # data is already in random order here, so use that order for deletion task
    for i in 1:15
        delete!(tree, data[i])
    end
    println("After deleting 15 of the nodes, ", tree)
end

RBTsort()
Output:
After adding 30 nodes, Red-Black tree node list is:
(left: empty, color: red, current: 12, right: empty, parent: 18)
(left: 12, color: black, current: 18, right: 19, parent: 58)
(left: empty, color: red, current: 19, right: empty, parent: 18)
(left: 18, color: red, current: 58, right: 67, parent: 94)
(left: empty, color: black, current: 67, right: 78, parent: 58)
(left: empty, color: red, current: 78, right: empty, parent: 67)
(left: 58, color: black, current: 94, right: 155, parent: 207)
(left: empty, color: black, current: 112, right: 153, parent: 155)
(left: empty, color: red, current: 153, right: empty, parent: 112)
(left: 112, color: red, current: 155, right: 187, parent: 94)
(left: empty, color: red, current: 181, right: empty, parent: 187)
(left: 181, color: black, current: 187, right: empty, parent: 155)
(left: 94, color: black, current: 207, right: 488, parent: empty)
(left: empty, color: black, current: 215, right: empty, parent: 253)
(left: 215, color: black, current: 253, right: 401, parent: 488)
(left: empty, color: red, current: 305, right: empty, parent: 319)
(left: 305, color: black, current: 319, right: 390, parent: 401)
(left: empty, color: red, current: 390, right: empty, parent: 319)
(left: 319, color: red, current: 401, right: 408, parent: 253)
(left: empty, color: black, current: 408, right: 445, parent: 401)
(left: empty, color: red, current: 445, right: empty, parent: 408)
(left: 253, color: red, current: 488, right: 862, parent: 207)
(left: empty, color: red, current: 606, right: empty, parent: 653)
(left: 606, color: black, current: 653, right: 731, parent: 754)
(left: empty, color: red, current: 731, right: empty, parent: 653)
(left: 653, color: red, current: 754, right: 842, parent: 862)
(left: empty, color: black, current: 842, right: empty, parent: 754)
(left: 754, color: black, current: 862, right: 1000, parent: 488)
(left: empty, color: red, current: 948, right: empty, parent: 1000)
(left: 948, color: black, current: 1000, right: empty, parent: 862)

After deleting 15 of the nodes, Red-Black tree node list is:
(left: empty, color: red, current: 12, right: empty, parent: 18)
(left: 12, color: black, current: 18, right: 19, parent: 58)
(left: empty, color: red, current: 19, right: empty, parent: 18)
(left: 18, color: red, current: 58, right: 78, parent: 153)
(left: empty, color: black, current: 78, right: empty, parent: 58)
(left: 58, color: black, current: 153, right: 181, parent: 305)
(left: empty, color: black, current: 181, right: empty, parent: 153)
(left: 153, color: black, current: 305, right: 606, parent: empty)
(left: empty, color: black, current: 319, right: empty, parent: 401)
(left: 319, color: black, current: 401, right: 445, parent: 606)
(left: empty, color: black, current: 445, right: empty, parent: 401)
(left: 401, color: red, current: 606, right: 754, parent: 305)
(left: empty, color: black, current: 731, right: empty, parent: 754)
(left: 731, color: black, current: 754, right: 948, parent: 606)
(left: empty, color: black, current: 948, right: empty, parent: 754)