去年做的轮子:AVL

还是去年做的轮子

module AVL

include("List.jl")
import .LinkedList
import .LinkedList:push
import .LinkedList:Queue
import .LinkedList:remove
#----------------#
#    AVLNode     #
#----------------#
mutable struct AVLNode{T}
    value::T
    left::Union{Nothing,AVLNode{T}}
    right::Union{Nothing,AVLNode{T}}
    height::Int
end

AVLNode{T}(val::T) where T = AVLNode{T}(val,nothing,nothing,0)

#---------------#
#     macro     #
#---------------#
macro valueof(node)
    quote
        $node.value
    end |> esc
end
#---------------#
#    height     #
#---------------#
function height(node::AVLNode{T}) where T
    return node.height
end

function height(node::Nothing)
    return -1
end


#---------------#
#      LL       #
#  rotate right #
#---------------#
function rotateLL(node::AVLNode{T}) where T
    left_node=node.left
    node.left=left_node.right
    left_node.right=node
    # update height

    node.height=max(height(node.left),height(node.right))+1
    left_node.height=max(height(left_node.left),height(left_node.right))+1

    return left_node
end


#--------------#
#      RR      #
#  rotate left #
#--------------#
function rotateRR(node::AVLNode{T}) where T
    right_node=node.right
    node.right=right_node.left
    right_node.left=node

    node.height=max(height(node.left),height(node.right))+1
    right_node.height=max(height(right_node.left),height(right_node.right))+1

    return right_node
end


#-------------#
#      LR     #
#rotate left right#
#-------------#
function rotateLR(node::AVLNode{T}) where T
    node.left=rotateRR(node.left)
    return rotateLL(node)
end


#--------------#
#     RL       #
#rotate right left#
#--------------#
function rotateRL(node::AVLNode{T}) where T
    node.right=rotateLL(node.right)
    return rotateRR(node)
end


#--------------#
#    insert    #
#--------------#
function insert_node(node::Union{Nothing,AVLNode{T}},data::T) where T
    if node==nothing
        node=AVLNode{T}(data)
    else
        # insert at left
        if data < node.value
            node.left=insert_node(node.left,data)

            if height(node.left)-height(node.right)==2
                if data < node.left.value
                    node=rotateLL(node)
                else
                    node=rotateLR(node)
                end
            end
        else
            # insert at right
            node.right=insert_node(node.right,data)

            if height(node.right)-height(node.left)==2
                if data >= node.right.value
                    node=rotateRR(node)
                else
                    node=rotateRL(node)
                end
            end
        end

    end
    #update height
    node.height=max(height(node.left),height(node.right))+1
    return node

end



#---------------#
#   AVLTree     #
#---------------#
mutable struct AVLTree{T}
    root::Union{Nothing,AVLNode{T}}
end

AVLTree{T}() where T = AVLTree{T}(nothing)


#_--------------#
#     push      #
#---------------#
function push(tree::AVLTree{T},data::T) where T
    tree.root=insert_node(tree.root,data)
end


#--------------#
#   search     #
#--------------#
function search(tree::AVLTree{T},data::T) where T
    current=tree.root
    while current != nothing
        if data < @valueof current
            current=current.left
        elseif data > @valueof current
            current=current.right
        else
            break
        end
    end
    return current
end        
#--------------#
#    print     #
#--------------#

function Base.print(tree::AVLTree{T}) where T
    q=Queue{AVLNode{T}}()
    push(q,tree.root)
    while q.size != 0
        current=remove(q)
        print(current.value,'\t')

        if current.left != nothing
            push(q,current.left)
        end

        if current.right != nothing
            push(q,current.right)
        end

    end

end

#--------------#
#    export    #
#--------------#
export AVLTree
export print
export push
export search
end
2 个赞