implement rank, to complete order-statistics using AVLTree

Author: eulerkochy

docs/src/avl_tree.md

@@ -27,6 +27,9 @@ julia> for k in 1:2:10
 
 julia> haskey(tree, 5)
 false
+
+julia> rank(tree, 17)
+4
 ```
 
 ```@meta

src/DataStructures.jl

@@ -60,7 +60,7 @@ module DataStructures
     export DiBitVector
 
     export RBTree, search_node, minimum_node
-    export AVLTree
+    export AVLTree, rank
 
     export findkey
 

src/avl_tree.jl

@@ -228,6 +228,27 @@ function Base.delete!(tree::AVLTree{K}, d::K) where K
     return tree
 end
 
+"""
+    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 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)

test/test_avl_tree.jl

@@ -131,4 +131,15 @@
             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 rank(t9, i) == i
+        end
+        @test_throws KeyError rank(t9, 21)
+    end
 end