Add functions, modify testset, address review

Author: eulerkochy

src/DataStructures.jl

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

src/red_black_tree.jl

@@ -1,7 +1,7 @@
 # it has unique keys
 # leftChild has keys which are less than the node
 # rightChild has keys which are greater than the node
-# color is true if it's a Red Node, else it's false
+# color is true if its a Red Node, else its false
 mutable struct RBTreeNode{K}
     color::Bool 
     data::Union{K, Nothing}
@@ -17,14 +17,6 @@ end
 RBTreeNode() = RBTreeNode{Any}()
 RBTreeNode(d) = RBTreeNode{Any}(d)
 
-function Base.isequal(x::RBTreeNode{K}, y::RBTreeNode{K})where K
-   return (x.color == y.color && x.data == y.data)
-end
-
-Base.:(==)(x::RBTreeNode{K}, y::RBTreeNode{K}) where K = isequal(x, y)
-
-Base.hash(x::RBTreeNode) = hash(x.data, hash(x.color))
-
 function create_null_node(K::Type)
     node = RBTreeNode{K}()
     node.color = false
@@ -34,14 +26,21 @@ end
 mutable struct RBTree{K}
     root::RBTreeNode{K}
     nil::RBTreeNode{K}
+    count::Int
 
     function RBTree{K}() where K 
-        new{K}(create_null_node(K), create_null_node(K))
+        rb = new()
+        rb.nil = create_null_node(K)
+        rb.root = Ref(rb.nil)[]
+        rb.count = 0
+        return rb
     end
 end
 
 RBTree() = RBTree{Any}()
 
+Base.length(tree::RBTree) = tree.count
+
 """
     search_node(tree, key)
 
@@ -62,13 +61,13 @@ function search_node(tree::RBTree{K}, d::K) where K
 end
 
 """
-    search_key(tree, key)
+    haskey(tree, key)
 
 Returns true if `key` is present in the `tree`, else returns false.    
 """
-search_key(tree, key)
+haskey(tree, key)
 
-function search_key(tree::RBTree{K}, d::K) where K 
+function Base.haskey(tree::RBTree{K}, d::K) where K 
     node = search_node(tree, d)
     return (node.data == d)
 end
@@ -171,11 +170,11 @@ function fix_insert!(tree::RBTree, node::RBTreeNode)
                 uncle.color = false
                 node = grand_parent
             else  # uncle is black in color
-                if (node == parent.rightChild) # node is rightChild of it's parent
+                if (node == parent.rightChild) # node is rightChild of its parent
                     node = parent
                     left_rotate!(tree, node)
                 end
-                # node is leftChild of it's parent
+                # node is leftChild of its parent
                 node.parent.color = false
                 node.parent.parent.color = true
                 right_rotate!(tree, node.parent.parent)
@@ -189,11 +188,11 @@ function fix_insert!(tree::RBTree, node::RBTreeNode)
                 uncle.color = false
                 node = grand_parent
             else  # uncle is black in color
-                if (node == parent.leftChild) # node is leftChild of it's parent
+                if (node == parent.leftChild) # node is leftChild of its parent
                     node = parent
                     right_rotate!(tree, node)
                 end
-                # node is rightChild of it's parent
+                # node is rightChild of its parent
                 node.parent.color = false
                 node.parent.parent.color = true
                 left_rotate!(tree, node.parent.parent)
@@ -212,8 +211,8 @@ insert!(tree, key)
 
 function Base.insert!(tree::RBTree{K}, d::K) where K
     # if the key exists in the tree, no need to insert 
-    search_key(tree, d) && return tree
-    # search_key(tree, d) && return tree
+    haskey(tree, d) && return tree
+
     # insert, if not present in the tree
     node = RBTreeNode{K}(d)
     node.leftChild = node.rightChild = tree.nil
@@ -227,9 +226,20 @@ function Base.insert!(tree::RBTree{K}, d::K) where K
     else
         fix_insert!(tree, node)
     end
+    tree.count += 1
     return tree
 end
 
+"""
+    insert!(tree, key)
+
+Inserts `key` in the `tree` if it is not present.
+"""
+function Base.push!(tree::RBTree{K}, key0) where K
+    key = convert(K, key0)
+    insert!(tree, key)
+end
+
 """
     delete_fix(tree::RBTree, node::Union{RBTreeNode, Nothing})
 
@@ -317,6 +327,7 @@ end
 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
@@ -326,7 +337,7 @@ end
 """
     delete!(tree::RBTree, key)
 
-Deletes `key` from `tree`, if present, else throws a KeyError.
+Deletes `key` from `tree`, if present, else returns the unmodified tree.
 """
 delete!(tree, key)
 
@@ -346,7 +357,7 @@ function Base.delete!(tree::RBTree{K}, d::K) where K
         end
     end
 
-    (z == tree.nil) && throw(KeyError(d))
+    (z == tree.nil) && return tree
 
     y = z
     y_original_color = y.color
@@ -377,4 +388,30 @@ function Base.delete!(tree::RBTree{K}, d::K) where K
     end
 
     !y_original_color && delete_fix(tree, x)
-end
+    tree.count -= 1
+    return tree
+end
+
+Base.in(key, tree::RBTree) = haskey(tree, key)
+
+"""
+    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::RBTree{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::RBTreeNode{K}) where K
+        if (node != tree.nil)
+            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_red_black_tree.jl

@@ -1,77 +1,137 @@
 include("../src/red_black_tree.jl")
 @testset "RBTree" begin
-    t = RBTree{Int}()
-    for i = 1:10000
-        insert!(t, i)
-    end
-    for i = 1:10000
-        @test search_key(t, i)
-    end
+    @testset "inserting values" begin
+        t = RBTree{Int}()
+        for i in 1:100
+            insert!(t, i)
+        end
 
-    for i = 10001:20000
-        @test !search_key(t, i)
-    end
+        @test length(t) == 100
 
-    for i = 1:2:10000
-        delete!(t, i)
-    end
+        for i in 1:100
+            @test haskey(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)
+        for i = 101:200
+            @test !haskey(t, i)
         end
     end
 
-    for i = 1:2:1000
-        insert!(t, i)
+    @testset "deleting values" begin
+        t = RBTree{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
 
-    for i = 1:1000
-        @test search_key(t, i)
+    @testset "handling different cases of delete!" begin
+        t2 = RBTree()
+        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
 
-    # test hash
-    for i = 1:1000
-        node = search_node(t, i)
-        @test hash(node) == hash(i, hash(node.color))
+    @testset "handling different cases of insert!" begin
+        nums = rand(1:100000, 1000)
+        t3 = RBTree()
+        uniq_nums = Set(nums)
+        for num in nums
+            insert!(t3, num)
+        end
+        @test length(t3) == length(uniq_nums)
     end
 
-    # for handling cases related to delete!
-    t2 = RBTree()
-    for i = 1:100000
-        insert!(t2, i)
+    @testset "in" begin
+        t4 = RBTree{Char}()
+        push!(t4, 'a')
+        push!(t4, 'b')
+        @test length(t4) == 2
+        @test in('a', t4)
+        @test !in('c', t4)
     end
 
-    nums = rand(1:100000, 1000)
-    visited = Set()
-    for num in nums 
-        if num in visited
-            @test_throws KeyError delete!(t2, num)
-        else
-            delete!(t2, num)
-            push!(visited, num) 
+    @testset "search_node" begin 
+        t5 = RBTree()
+        for i in 1:32
+            push!(t5, i)
         end
-    end
+        n1 = search_node(t5, 21)
+        @test n1.data == 21
+        n2 = search_node(t5, 35)
+        @test n2 === t5.nil
+        n3 = search_node(t5, 0)
+        @test n3 === t5.nil
+    end 
 
-    for i = 1:100000
-        if i in visited
-            @test !search_key(t2, i)
-        else
-            @test search_key(t2, i)
+    @testset "getindex" begin 
+        t6 = RBTree{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
 
-    # for handling cases related to insert!
-    t3 = RBTree()
-    for num in nums
-        insert!(t3, num)
+    @testset "key conversion in push!" begin
+        t7 = RBTree{Int}()
+        push!(t7, Int8(1))
+        @test length(t7) == 1
+        @test haskey(t7, 1)
     end
 
-    for i in visited
-        @test search_key(t3, i)
-    end
+    @testset "minimum_node" begin
+        t8 = RBTree()
+        for i in 1:32
+            push!(t8, i)
+        end
+        m1 = minimum_node(t8, t8.root)
+        @test m1.data == 1
+        node = t8.root
+        while node.leftChild != t8.nil
+            m = minimum_node(t8, node.leftChild)
+            @test m == m1
+            node = node.leftChild
+        end
 
+        @test minimum_node(t8, t8.nil) === t8.nil
+    end
 end