increase coverage, add tests

Author: eulerkochy

src/avl_tree.jl

@@ -1,10 +1,13 @@
+# 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}
     data::K
 
-    AVLTreeNode{K}() where K = new{K}(0, nothing, nothing)
     AVLTreeNode{K}(d::K) where K = new{K}(1, nothing, nothing, d)
 end
 
@@ -37,7 +40,7 @@ end
 fix_height(node::AVLTreeNode) = 1 + max(get_height(node.leftChild), get_height(node.rightChild))
 
 """
-    left_rotate!(node_x::RBTreeNode)
+    left_rotate!(node_x::AVLTreeNode)
 
 Performs a left-rotation on `node_x`, updates height of the nodes, and returns the rotated node. 
 """
@@ -52,7 +55,7 @@ function left_rotate(z::AVLTreeNode)
 end
 
 """
-    right_rotate!(node_x::RBTreeNode)
+    right_rotate!(node_x::AVLTreeNode)
 
 Performs a right-rotation on `node_x`, updates height of the nodes, and returns the rotated node. 
 """
@@ -67,9 +70,9 @@ function right_rotate(z::AVLTreeNode)
 end
 
 """
-   minimum_node(tree::RBTree, node::RBTreeNode) 
+   minimum_node(tree::AVLTree, node::AVLTreeNode) 
 
-Returns the RBTreeNode with minimum value in subtree of `node`. 
+Returns the AVLTreeNode with minimum value in subtree of `node`. 
 """
 function minimum_node(node::Union{AVLTreeNode, Nothing})
     while node != nothing && node.leftChild != nothing
@@ -114,6 +117,8 @@ function haskey(tree::AVLTree{K}, d::K) where K
     return (node.data == d)
 end
 
+Base.in(key, tree::AVLTree) = haskey(tree, key)
+
 """
     insert!(tree, key)
 
@@ -138,7 +143,7 @@ function Base.insert!(tree::AVLTree{K}, d::K) where K
         if balance > 1
             if key < node.leftChild.data
                 return right_rotate(node)
-
+            else
                 node.leftChild = left_rotate(node.leftChild)
                 return right_rotate(node)
             end
@@ -198,10 +203,6 @@ function Base.delete!(tree::AVLTree{K}, d::K) where K
             end 
         end
 
-        if node == nothing
-            return node
-        end
-
         node.height = fix_height(node)
         balance = get_balance(node)
 
@@ -234,3 +235,25 @@ function Base.delete!(tree::AVLTree{K}, d::K) where K
     tree.count -= 1
     return tree
 end
+
+"""
+    getindex(tree, ind)
+
+Gets the key present at index `ind` of the tree. Indexing is done in increasing order of key.
+"""
+getindex(tree, ind)
+
+function Base.getindex(tree::AVLTree{K}, ind) where K 
+    @boundscheck (1 <= ind <= tree.count) || throw(ArgumentError("$ind should be in between 1 and $(tree.count)"))
+    function traverse_tree_inorder(node::Union{AVLTreeNode, Nothing})
+        if (node != nothing)
+            left = traverse_tree_inorder(node.leftChild)
+            right = traverse_tree_inorder(node.rightChild)
+            append!(push!(left, node.data), right)
+        else
+            return K[]
+        end
+    end
+    arr = traverse_tree_inorder(tree.root) 
+    return @inbounds arr[ind]
+end

test/test_avl_tree.jl

@@ -1,7 +1,7 @@
 @testset "AVLTree" begin
     @testset "inserting values" begin
         t = AVLTree{Int}()
-        for i in 1:100
+        for i in 100:-1:1
             insert!(t, i)
         end
 
@@ -11,7 +11,7 @@
             @test haskey(t, i)
         end
 
-        for i = 101:200
+        for i = 101:150
             @test !haskey(t, i)
         end
     end
@@ -75,4 +75,60 @@
         @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 ArgumentError getindex(t6, 0)
+        @test_throws ArgumentError 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
 end