Repository: JuliaCollections/DataStructures.jl
Pull request #640 (2020-07-06 16:35:29)

AVL Tree
Part of GSoC'20 evaluation
This PR implements AVL Tree. 


Patch:
diff --git a/README.md b/README.md
index 2d6f999..03429ea 100644
--- a/README.md
+++ b/README.md
@@ -31,6 +31,7 @@ This package implements a variety of data structures, including
 -   SparseIntSet
 -   DiBitVector (in which each element can store two bits)
 -   Red Black Tree
+-   AVL Tree
 
 Resources
 ---------
diff --git a/docs/make.jl b/docs/make.jl
index cca1323..56e9c88 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -26,6 +26,7 @@ makedocs(
         "sorted_containers.md",
         "dibit_vector.md",
         "red_black_tree.md",
+        "avl_tree.md",
     ],
     modules = [DataStructures],
     format = Documenter.HTML()
diff --git a/docs/src/avl_tree.md b/docs/src/avl_tree.md
new file mode 100644
index 0000000..3256654
--- /dev/null
+++ b/docs/src/avl_tree.md
@@ -0,0 +1,37 @@
+```@meta
+DocTestSetup = :(using DataStructures)
+```
+
+# AVL Tree
+
+The `AVLTree` type is an implementation of AVL Tree in Julia. It is a self-balancing binary search tree where balancing occurs based on the difference of height of the left subtree and the right subtree. Operations such as search, insert and delete can be done in `O(log n)` complexity, where `n` is the number of nodes in the `AVLTree`. Order-statistics on the keys can also be done in `O(log n)`.
+
+Examples:
+
+```jldoctest
+julia> tree = AVLTree{Int}();
+
+julia> for k in 1:2:20
+           push!(tree, k)
+       end
+
+julia> haskey(tree, 3)
+true
+
+julia> tree[4] # time complexity of this operation is O(log n)
+7
+
+julia> for k in 1:2:10
+           delete!(tree, k)
+       end
+
+julia> haskey(tree, 5)
+false
+
+julia> sorted_rank(tree, 17) # used for finding rank of the key
+4
+```
+
+```@meta
+DocTestSetup = nothing
+```
diff --git a/docs/src/index.md b/docs/src/index.md
index 6577fa3..1818935 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -23,6 +23,7 @@ This package implements a variety of data structures, including
 -   SparseIntSet
 -   DiBitVector
 -   Red Black Tree
+-   AVL Tree
 
 ## Contents
 
@@ -48,6 +49,7 @@ Pages = [
     "sorted_containers.md",
     "sparse_int_set.md",
     "dibit_vector.md",
-    "red_black_tree.md"
+    "red_black_tree.md",
+    "avl_tree.md"
 ]
 ```
diff --git a/src/DataStructures.jl b/src/DataStructures.jl
index acc0dd7..675c625 100644
--- a/src/DataStructures.jl
+++ b/src/DataStructures.jl
@@ -60,6 +60,7 @@ module DataStructures
     export DiBitVector
 
     export RBTree, search_node, minimum_node
+    export AVLTree, sorted_rank
 
     export findkey
 
@@ -112,7 +113,8 @@ module DataStructures
     export SparseIntSet
 
     include("dibit_vector.jl")
+    include("avl_tree.jl")
     include("red_black_tree.jl")
-    include("deprecations.jl")
 
+    include("deprecations.jl")
 end
diff --git a/src/avl_tree.jl b/src/avl_tree.jl
new file mode 100644
index 0000000..bacf1cb
--- /dev/null
+++ b/src/avl_tree.jl
@@ -0,0 +1,268 @@
+# it has unique keys
+# leftChild has keys which are less than the node
+# rightChild has keys which are greater than the node
+# height stores the height of the subtree.
+mutable struct AVLTreeNode{K}
+    height::Int8
+    leftChild::Union{AVLTreeNode{K}, Nothing}
+    rightChild::Union{AVLTreeNode{K}, Nothing}
+    subsize::Int32
+    data::K
+
+    AVLTreeNode{K}(d::K) where K = new{K}(1, nothing, nothing, 1, d)
+end
+
+AVLTreeNode(d) = AVLTreeNode{Any}(d)
+
+AVLTreeNode_or_null{T} = Union{AVLTreeNode{T}, Nothing}
+
+mutable struct AVLTree{T}
+    root::AVLTreeNode_or_null{T}
+    count::Int
+
+    AVLTree{T}() where T = new{T}(nothing, 0)
+end
+
+AVLTree() = AVLTree{Any}()
+
+Base.length(tree::AVLTree) = tree.count
+
+get_height(node::Union{AVLTreeNode, Nothing}) = (node == nothing) ? 0 : node.height
+
+# balance is the difference of height between leftChild and rightChild of a node.
+function get_balance(node::Union{AVLTreeNode, Nothing})
+    if node == nothing
+        return 0
+    else
+        return get_height(node.leftChild) - get_height(node.rightChild)
+    end
+end
+
+# computes the height of the subtree, which basically is
+# one added the maximum of the height of the left subtree and right subtree
+compute_height(node::AVLTreeNode) = 1 + max(get_height(node.leftChild), get_height(node.rightChild))
+
+get_subsize(node::AVLTreeNode_or_null) = (node == nothing) ? 0 : node.subsize
+
+# compute the subtree size
+function compute_subtree_size(node::AVLTreeNode_or_null)
+    if node == nothing
+        return 0
+    else
+        L = get_subsize(node.leftChild)
+        R = get_subsize(node.rightChild)
+        return (L + R + 1)
+    end
+end
+
+"""
+    left_rotate(node_x::AVLTreeNode)
+
+Performs a left-rotation on `node_x`, updates height of the nodes, and returns the rotated node. 
+"""
+function left_rotate(z::AVLTreeNode)
+    y = z.rightChild
+    α = y.leftChild
+    y.leftChild = z
+    z.rightChild = α
+    z.height = compute_height(z)
+    y.height = compute_height(y)
+    z.subsize = compute_subtree_size(z)
+    y.subsize = compute_subtree_size(y)
+    return y
+end
+
+"""
+    right_rotate(node_x::AVLTreeNode)
+
+Performs a right-rotation on `node_x`, updates height of the nodes, and returns the rotated node. 
+"""
+function right_rotate(z::AVLTreeNode)
+    y = z.leftChild
+    α = y.rightChild
+    y.rightChild = z
+    z.leftChild = α
+    z.height = compute_height(z)
+    y.height = compute_height(y)
+    z.subsize = compute_subtree_size(z)
+    y.subsize = compute_subtree_size(y)
+    return y
+end
+
+"""
+   minimum_node(tree::AVLTree, node::AVLTreeNode) 
+
+Returns the AVLTreeNode with minimum value in subtree of `node`. 
+"""
+function minimum_node(node::Union{AVLTreeNode, Nothing})
+    while node != nothing && node.leftChild != nothing
+        node = node.leftChild
+    end
+    return node
+end
+
+function search_node(tree::AVLTree{K}, d::K) where K
+    prev = nothing
+    node = tree.root
+    while node != nothing && node.data != nothing && node.data != d
+
+        prev = node
+        if d < node.data 
+            node = node.leftChild
+        else
+            node = node.rightChild
+        end
+    end
+    
+    return (node == nothing) ? prev : node
+end
+
+function Base.haskey(tree::AVLTree{K}, d::K) where K 
+    (tree.root == nothing) && return false
+    node = search_node(tree, d)
+    return (node.data == d)
+end
+
+Base.in(key, tree::AVLTree) = haskey(tree, key)
+
+function Base.insert!(tree::AVLTree{K}, d::K) where K
+
+    function insert_node(node::Union{AVLTreeNode, Nothing}, key)
+        if node == nothing
+            return AVLTreeNode{K}(key)
+        elseif key < node.data
+            node.leftChild = insert_node(node.leftChild, key)
+        else
+            node.rightChild = insert_node(node.rightChild, key)
+        end
+        
+        node.subsize = compute_subtree_size(node)
+        node.height = compute_height(node)
+        balance = get_balance(node)
+        
+        if balance > 1
+            if key < node.leftChild.data
+                return right_rotate(node)
+            else
+                node.leftChild = left_rotate(node.leftChild)
+                return right_rotate(node)
+            end
+        end
+
+        if balance < -1
+            if key > node.rightChild.data
+                return left_rotate(node)
+            else
+                node.rightChild = right_rotate(node.rightChild)
+                return left_rotate(node)
+            end
+        end
+
+        return node
+    end
+
+    haskey(tree, d) && return tree
+
+    tree.root = insert_node(tree.root, d)
+    tree.count += 1
+    return tree
+end
+
+function Base.push!(tree::AVLTree{K}, key0) where K
+    key = convert(K, key0)
+    insert!(tree, key)
+end
+
+function Base.delete!(tree::AVLTree{K}, d::K) where K
+
+    function delete_node!(node::Union{AVLTreeNode, Nothing}, key)
+        if key < node.data
+            node.leftChild = delete_node!(node.leftChild, key)
+        elseif key > node.data
+            node.rightChild = delete_node!(node.rightChild, key)
+        else
+            if node.leftChild == nothing
+                result = node.rightChild
+                return result
+            elseif node.rightChild == nothing
+                result = node.leftChild
+                return result
+            else 
+                result = minimum_node(node.rightChild)
+                node.data = result.data
+                node.rightChild = delete_node!(node.rightChild, result.data)
+            end 
+        end
+        
+        node.subsize = compute_subtree_size(node)
+        node.height = compute_height(node)
+        balance = get_balance(node)
+
+        if balance > 1
+            if get_balance(node.leftChild) >= 0
+                return right_rotate(node)
+            else
+                node.leftChild = left_rotate(node.leftChild)
+                return right_rotate(node)
+            end
+        end
+
+        if balance < -1
+            if get_balance(node.rightChild) <= 0
+                return left_rotate(node)
+            else
+                node.rightChild = right_rotate(node.rightChild)
+                return left_rotate(node)
+            end
+        end 
+        
+        return node
+    end
+
+    # if the key is not in the tree, do nothing and return the tree
+    !haskey(tree, d) && return tree
+    
+    # if the key is present, delete it from the tree
+    tree.root = delete_node!(tree.root, d)
+    tree.count -= 1
+    return tree
+end
+
+"""
+    sorted_rank(tree::AVLTree, key)
+
+Returns the rank of `key` present in the `tree`, if it present. A KeyError is thrown if `key` is not present.
+"""
+function sorted_rank(tree::AVLTree{K}, key::K) where K
+    !haskey(tree, key) && throw(KeyError(key))
+    node = tree.root
+    rank = 0
+    while node.data != key
+        if (node.data < key)
+            rank += (1 + get_subsize(node.leftChild))
+            node = node.rightChild
+        else
+            node = node.leftChild
+        end
+    end 
+    rank += (1 + get_subsize(node.leftChild))
+    return rank
+end
+
+function Base.getindex(tree::AVLTree{K}, ind::Integer) where K 
+    @boundscheck (1 <= ind <= tree.count) || throw(BoundsError("$ind should be in between 1 and $(tree.count)"))
+    function traverse_tree(node::AVLTreeNode_or_null, idx)
+        if (node != nothing)
+            L = get_subsize(node.leftChild)
+            if idx <= L
+                return traverse_tree(node.leftChild, idx)
+            elseif idx == L + 1
+                return node.data
+            else
+                return traverse_tree(node.rightChild, idx - L - 1)
+            end
+        end
+    end
+    value = traverse_tree(tree.root, ind) 
+    return value
+end
\ No newline at end of file


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 7cce031..f81b35a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -32,7 +32,8 @@ tests = ["deprecations",
          "robin_dict",
          "ordered_robin_dict",
          "dibit_vector",
-         "red_black_tree"
+         "red_black_tree",
+         "avl_tree"
         ]
 
 if length(ARGS) > 0
diff --git a/test/test_avl_tree.jl b/test/test_avl_tree.jl
new file mode 100644
index 0000000..78ddae4
--- /dev/null
+++ b/test/test_avl_tree.jl
@@ -0,0 +1,145 @@
+@testset "AVLTree" begin
+    @testset "inserting values" begin
+        t = AVLTree{Int}()
+        for i in 100:-1:1
+            insert!(t, i)
+        end
+
+        @test length(t) == 100
+
+        for i in 1:100
+            @test haskey(t, i)
+        end
+
+        for i = 101:150
+            @test !haskey(t, i)
+        end
+    end
+
+    @testset "deleting values" begin
+        t = AVLTree{Int}()
+        for i in 1:100
+            insert!(t, i)
+        end
+        for i in 1:2:100
+            delete!(t, i)
+        end
+
+        @test length(t) == 50
+        
+        for i in 1:100
+            if iseven(i)
+                @test haskey(t, i)
+            else
+                @test !haskey(t, i)
+            end
+        end
+
+        for i in 1:2:100
+            insert!(t, i)
+        end
+
+        @test length(t) == 100
+    end
+
+    @testset "handling different cases of delete!" begin
+        t2 = AVLTree()
+        for i in 1:100000
+            insert!(t2, i)
+        end
+
+        @test length(t2) == 100000
+        
+        nums = rand(1:100000, 8599)
+        visited = Set()
+        for num in nums 
+            if !(num in visited)
+                delete!(t2, num)
+                push!(visited, num) 
+            end
+        end
+
+        for i in visited
+            @test !haskey(t2, i)
+        end
+        @test (length(t2) + length(visited)) == 100000
+    end
+
+    @testset "handling different cases of insert!" begin
+        nums = rand(1:100000, 1000)
+        t3 = AVLTree()
+        uniq_nums = Set(nums)
+        for num in uniq_nums
+            insert!(t3, num)
+        end
+        @test length(t3) == length(uniq_nums)
+    end
+
+    @testset "in" begin
+        t4 = AVLTree{Char}()
+        push!(t4, 'a')
+        push!(t4, 'b')
+        @test length(t4) == 2
+        @test in('a', t4)
+        @test !in('c', t4)
+    end
+
+    @testset "search_node" begin 
+        t5 = AVLTree()
+        for i in 1:32
+            push!(t5, i)
+        end
+        n1 = search_node(t5, 21)
+        @test n1.data == 21
+        n2 = search_node(t5, 35)
+        @test n2.data == 32 
+        n3 = search_node(t5, 0)
+        @test n3.data == 1
+    end 
+
+    @testset "getindex" begin 
+        t6 = AVLTree{Int}()
+        for i in 1:10
+            push!(t6, i)
+        end
+        for i in 1:10
+            @test getindex(t6, i) == i
+        end
+        @test_throws BoundsError getindex(t6, 0)
+        @test_throws BoundsError getindex(t6, 11)
+    end
+
+    @testset "key conversion in push!" begin
+        t7 = AVLTree{Int}()
+        push!(t7, Int8(1))
+        @test length(t7) == 1
+        @test haskey(t7, 1)
+    end
+
+    @testset "minimum_node" begin
+        t8 = AVLTree()
+        @test minimum_node(t8.root) == nothing
+        for i in 1:32
+            push!(t8, i)
+        end
+        m1 = minimum_node(t8.root)
+        @test m1.data == 1
+        node = t8.root
+        while node.leftChild != nothing
+            m = minimum_node(node.leftChild)
+            @test m == m1
+            node = node.leftChild
+        end
+    end
+
+    @testset "rank" begin 
+        t9 = AVLTree{Int}()
+        for i in 1:20
+            push!(t9, i)
+        end
+        for i in 1:20
+            @test sorted_rank(t9, i) == i
+        end
+        @test_throws KeyError sorted_rank(t9, 21)
+    end
+end
\ No newline at end of file