add search_key, add tests

Author: eulerkochy

src/DataStructures.jl

@@ -110,7 +110,9 @@ module DataStructures
     export SparseIntSet
 
     include("dibit_vector.jl")
+    include("avl_tree.jl")
     include("red_black_tree.jl")
+
     include("deprecations.jl")
 
 end

src/avl_tree.jl

@@ -1,11 +1,11 @@
 mutable struct AVLTreeNode{K}
     height::Int8
-    data::K
     leftChild::Union{AVLTreeNode{K}, Nothing}
     rightChild::Union{AVLTreeNode{K}, Nothing}
+    data::K
 
-    AVLTreeNode{K}() where K = new{K}(0, nothing, nothing, nothing)
-    AVLTreeNode{K}(d::K) where K = new{K}(1, d, nothing, nothing)
+    AVLTreeNode{K}() where K = new{K}(0, nothing, nothing)
+    AVLTreeNode{K}(d::K) where K = new{K}(1, nothing, nothing, d)
 end
 
 AVLTreeNode(d) = AVLTreeNode{Any}(d)
@@ -15,7 +15,7 @@ AVLTreeNode_or_null{T} = Union{AVLTreeNode{T}, Nothing}
 mutable struct AVLTree{T}
     root::AVLTreeNode_or_null{T}
 
-    AVLTree{T}() where T = new{T}(AVLTreeNode{T}())
+    AVLTree{T}() where T = new{T}(nothing)
 end
 
 AVLTree() = AVLTree{Any}()
@@ -65,6 +65,27 @@ function get_minimum_node(node::Union{AVLTreeNode, Nothing})
     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 search_key(tree::AVLTree{K}, d::K) where K 
+    (tree.root == nothing) && return false
+    node = search_node(tree, d)
+    return (node.data == d)
+end
+
 function Base.insert!(tree::AVLTree{K}, d::K) where K
 
     function insert_node(node::Union{AVLTreeNode, Nothing}, key)
@@ -99,6 +120,9 @@ function Base.insert!(tree::AVLTree{K}, d::K) where K
 
         return node
     end
+
+    search_key(tree, d) && return tree
+
     tree.root = insert_node(tree.root, d)
     return tree
 end
@@ -156,6 +180,8 @@ function Base.delete!(tree::AVLTree{K}, d::K) where K
         return node
     end
 
+    !search_key(tree, d) && throw(KeyError(d))
+    
     tree.root = delete_node(tree.root, d)
     return tree
 end

test/runtests.jl

@@ -31,7 +31,8 @@ tests = ["deprecations",
          "fenwick",
          "robin_dict",
          "dibit_vector",
-         "red_black_tree"
+         "red_black_tree",
+         "avl_tree",
         ]
 
 if length(ARGS) > 0

test/test_avl_tree.jl

@@ -0,0 +1,61 @@
+include("../src/avl_tree.jl")
+@testset "AVLTree" begin
+    t = AVLTree()
+    for i = 1:10000
+    	insert!(t, i)
+    end
+    for i = 1:10000
+    	@test search_key(t, i)
+    end
+
+    for i = 10001:20000
+        @test !search_key(t, i)
+    end
+
+    for i = 1:2:10000
+    	delete!(t, i)
+    end
+
+    for i = 1:10000
+    	if iseven(i)
+    		@test search_key(t, i)
+    	else
+    		@test !search_key(t, i)
+    		@test_throws KeyError delete!(t, i)
+    	end
+    end
+
+    for i = 1:2:1000
+        insert!(t, i)
+    end
+
+    for i = 1:1000
+        @test search_key(t, i)
+    end
+
+    # for handling cases related to delete!
+    t2 = AVLTree{Int}()
+    for i = 1:100000
+        insert!(t2, i)
+    end
+
+    nums = rand(1:100000, 1000)
+    visited = Set{Int}()
+    for num in nums 
+        if num in visited
+            @test_throws KeyError delete!(t2, num)
+        else
+            delete!(t2, num)
+            push!(visited, num) 
+        end
+    end
+
+    for i = 1:100000
+        if i in visited
+            @test !search_key(t2, i)
+        else
+            @test search_key(t2, i)
+        end
+    end
+
+end