Repository: JuliaCollections/DataStructures.jl
Pull request #645 (2020-07-22 15:26:36)

SplayTree
Part of GSoC'20 implementation.

SplayTree has no additional memory footprint as compared to Red-Black Tree or AVL Tree. It offers `O(log n)` amortized time complexity for the operations `search`, `insert` and `delete` 

TO-DO

- [x] writes tests


Patch:
diff --git a/README.md b/README.md
index 9e99bcc..d1409fb 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,7 @@ This package implements a variety of data structures, including
 -   DiBitVector (in which each element can store two bits)
 -   Red Black Tree
 -   AVL Tree
+-   Splay Tree
 
 Resources
 ---------
diff --git a/docs/make.jl b/docs/make.jl
index c75f845..ca04d66 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -28,6 +28,7 @@ makedocs(
         "dibit_vector.md",
         "red_black_tree.md",
         "avl_tree.md",
+        "splay_tree.md",
     ],
     modules = [DataStructures],
     format = Documenter.HTML()
diff --git a/docs/src/index.md b/docs/src/index.md
index 174bb74..416b07d 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -25,6 +25,7 @@ This package implements a variety of data structures, including
 -   DiBitVector
 -   Red Black Tree
 -   AVL Tree
+-   Splay Tree
 
 ## Contents
 
@@ -51,6 +52,7 @@ Pages = [
     "sparse_int_set.md",
     "dibit_vector.md",
     "red_black_tree.md",
-    "avl_tree.md"
+    "avl_tree.md",
+    "splay_tree.md",
 ]
 ```
diff --git a/docs/src/splay_tree.md b/docs/src/splay_tree.md
new file mode 100644
index 0000000..fb7fc82
--- /dev/null
+++ b/docs/src/splay_tree.md
@@ -0,0 +1,34 @@
+```@meta
+DocTestSetup = :(using DataStructures)
+```
+
+# Splay Tree
+
+The `SplayTree` type is an implementation of Splay Tree in Julia. It is a self-balancing binary search tree with the additional property that recently accessed elements are quick to access again. Operations such as search, insert and delete can be done in `O(log n)` amortized time, where `n` is the number of nodes in the `SplayTree`.
+
+Examples:
+
+```jldoctest
+julia> tree = SplayTree{Int}();
+
+julia> for k in 1:2:20
+           push!(tree, k)
+       end
+
+julia> haskey(tree, 3)
+true
+
+julia> tree[4]
+7
+
+julia> for k in 1:2:10
+           delete!(tree, k)
+       end
+
+julia> haskey(tree, 5)
+false
+```
+
+```@meta
+DocTestSetup = nothing
+```
diff --git a/src/DataStructures.jl b/src/DataStructures.jl
index ce37b06..3181475 100644
--- a/src/DataStructures.jl
+++ b/src/DataStructures.jl
@@ -51,7 +51,9 @@ module DataStructures
     export DiBitVector
 
     export RBTree, search_node, minimum_node
+    export SplayTree, maximum_node
     export AVLTree, sorted_rank
+    export SplayTree, maximum_node
 
     export findkey
 
@@ -107,6 +109,7 @@ module DataStructures
     include("dibit_vector.jl")
     include("avl_tree.jl")
     include("red_black_tree.jl")
+    include("splay_tree.jl")
 
     include("deprecations.jl")
 end
diff --git a/src/splay_tree.jl b/src/splay_tree.jl
new file mode 100644
index 0000000..a87614b
--- /dev/null
+++ b/src/splay_tree.jl
@@ -0,0 +1,227 @@
+mutable struct SplayTreeNode{K}
+    leftChild::Union{SplayTreeNode{K}, Nothing}
+    rightChild::Union{SplayTreeNode{K}, Nothing}
+    parent::Union{SplayTreeNode{K}, Nothing}
+    data::K
+
+    SplayTreeNode{K}() where K = new{K}(nothing, nothing, nothing)
+    SplayTreeNode{K}(d::K) where K = new{K}(nothing, nothing, nothing, d)
+end
+
+SplayTreeNode(d) = SplayTreeNode{Any}(d)
+SplayTreeNode() = SplayTreeNode{Any}()
+
+mutable struct SplayTree{K}
+    root::Union{SplayTreeNode{K}, Nothing}
+    count::Int
+
+    SplayTree{K}() where K = new{K}(nothing, 0)
+end 
+
+Base.length(tree::SplayTree) = tree.count
+
+SplayTree() = SplayTree{Any}()
+
+function left_rotate!(tree::SplayTree, node_x::SplayTreeNode)
+    node_y = node_x.rightChild
+    node_x.rightChild = node_y.leftChild
+    if node_y.leftChild != nothing
+        node_y.leftChild.parent = node_x
+    end
+    node_y.parent = node_x.parent
+
+    if node_x.parent == nothing
+        tree.root = node_y
+    elseif (node_x == node_x.parent.leftChild)
+        node_x.parent.leftChild = node_y
+    else
+        node_x.parent.rightChild = node_y
+    end
+    if node_y != nothing
+        node_y.leftChild = node_x
+    end
+    node_x.parent = node_y
+end    
+
+function right_rotate!(tree::SplayTree, node_x::SplayTreeNode)
+    node_y = node_x.leftChild
+    node_x.leftChild = node_y.rightChild
+    if node_y.rightChild != nothing
+        node_y.rightChild.parent = node_x
+    end
+    node_y.parent = node_x.parent
+    if node_x.parent == nothing
+        tree.root = node_y
+    elseif (node_x == node_x.parent.leftChild)
+        node_x.parent.leftChild = node_y
+    else
+        node_x.parent.rightChild = node_y
+    end
+    node_y.rightChild = node_x
+    node_x.parent = node_y
+end 
+
+# The splaying operation moves node_x to the root of the tree using the series of rotations.
+function splay!(tree::SplayTree, node_x::SplayTreeNode)
+    while node_x.parent !== nothing
+        parent = node_x.parent
+        grand_parent = node_x.parent.parent
+        if grand_parent === nothing
+            # single rotation
+            if node_x == parent.leftChild
+                # zig rotation
+                right_rotate!(tree, node_x.parent)
+            else 
+                # zag rotation
+                left_rotate!(tree, node_x.parent)
+            end
+            # double rotation
+        elseif node_x == parent.leftChild && parent == grand_parent.leftChild
+            # zig-zig rotation
+            right_rotate!(tree, grand_parent)
+            right_rotate!(tree, parent)
+        elseif node_x == parent.rightChild && parent == grand_parent.rightChild
+            # zag-zag rotation
+            left_rotate!(tree, grand_parent)
+            left_rotate!(tree, parent)
+        elseif node_x == parent.rightChild && parent == grand_parent.leftChild
+            # zig-zag rotation
+            left_rotate!(tree, node_x.parent)
+            right_rotate!(tree, node_x.parent)
+        else
+            # zag-zig rotation
+            right_rotate!(tree, node_x.parent)
+            left_rotate!(tree, node_x.parent)
+        end
+    end
+end
+
+function maximum_node(node::Union{SplayTreeNode, Nothing})
+    (node == nothing) && return node
+    while node.rightChild != nothing
+        node = node.rightChild
+    end
+    return node
+end
+
+# Join operations joins two trees S and T 
+# All the items in S are smaller than the items in T.
+# This is a two-step process.
+# In the first step, splay the largest node in S. This moves the largest node to the root node.
+# In the second step, set the right child of the new root of S to T.
+function _join!(tree::SplayTree, s::Union{SplayTreeNode, Nothing}, t::Union{SplayTreeNode, Nothing})
+    if s === nothing
+        return t
+    elseif t === nothing
+        return s
+    else
+        x = maximum_node(s)
+        splay!(tree, x)
+        x.rightChild = t
+        t.parent = x
+        return x
+    end
+end
+
+function search_node(tree::SplayTree{K}, d::K) where K
+    node = tree.root
+    prev = nothing
+    while node != nothing && node.data != d
+        prev = node
+        if node.data < d
+            node = node.rightChild
+        else
+            node = node.leftChild
+        end
+    end
+    return (node == nothing) ? prev : node
+end
+
+function Base.haskey(tree::SplayTree{K}, d::K) where K
+    node = tree.root
+    if node === nothing
+        return false
+    else
+        node = search_node(tree, d)
+        (node === nothing) && return false
+        is_found = (node.data == d)
+        is_found && splay!(tree, node)
+        return is_found
+    end
+end
+
+Base.in(key, tree::SplayTree) = haskey(tree, key)
+
+function Base.delete!(tree::SplayTree{K}, d::K) where K
+    node = tree.root
+    x = search_node(tree, d)
+    (x == nothing) && return tree
+    t = nothing
+    s = nothing 
+    
+    splay!(tree, x)
+    
+    if x.rightChild !== nothing
+        t = x.rightChild
+        t.parent = nothing
+    end
+
+    s = x
+    s.rightChild = nothing
+
+    if s.leftChild !== nothing
+        s.leftChild.parent = nothing
+    end
+
+    tree.root = _join!(tree, s.leftChild, t)
+    tree.count -= 1
+    return tree
+end
+
+function Base.push!(tree::SplayTree{K}, d0) where K
+    d = convert(K, d0)
+    is_present = search_node(tree, d)
+    if (is_present !== nothing) && (is_present.data == d)
+        return tree
+    end
+    # only unique keys are inserted
+    node = SplayTreeNode{K}(d)
+    y = nothing
+    x = tree.root
+
+    while x !== nothing
+        y = x
+        if node.data > x.data
+            x = x.rightChild
+        else
+            x = x.leftChild
+        end
+    end
+    node.parent = y
+
+    if y === nothing
+        tree.root = node
+    elseif node.data < y.data
+        y.leftChild = node
+    else
+        y.rightChild = node
+    end
+    splay!(tree, node)
+    tree.count += 1
+    return tree
+end
+
+function Base.getindex(tree::SplayTree{K}, ind) where K 
+    @boundscheck (1 <= ind <= tree.count) || throw(KeyError("$ind should be in between 1 and $(tree.count)"))
+    function traverse_tree_inorder(node::Union{SplayTreeNode, 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 patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 0daa7df..7a5aa76 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -32,9 +32,10 @@ tests = ["deprecations",
          "robin_dict",
          "ordered_robin_dict",
          "dibit_vector",
-         "red_black_tree",
          "swiss_dict",
-         "avl_tree"
+         "avl_tree",
+         "red_black_tree", 
+         "splay_tree"
         ]
 
 if length(ARGS) > 0
diff --git a/test/test_splay_tree.jl b/test/test_splay_tree.jl
new file mode 100644
index 0000000..bcd4041
--- /dev/null
+++ b/test/test_splay_tree.jl
@@ -0,0 +1,134 @@
+@testset "SplayTree" begin
+    @testset "pushing values" begin
+        t = SplayTree{Int}()
+        for i in 1:100
+            push!(t, i)
+        end
+
+        @test length(t) == 100
+
+        for i in 1:100
+            @test haskey(t, i)
+        end
+
+        for i = 101:200
+            @test !haskey(t, i)
+        end
+    end
+
+    @testset "deleting values" begin
+        t = SplayTree{Int}()
+        for i in 1:100
+            push!(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
+            push!(t, i)
+        end
+
+        @test length(t) == 100
+    end
+
+    @testset "handling different cases of delete!" begin
+        t2 = SplayTree()
+        for i in 1:100000
+            push!(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 push!" begin
+        nums = rand(1:100000, 1000)
+        t3 = SplayTree()
+        uniq_nums = Set(nums)
+        for num in nums
+            push!(t3, num)
+        end
+        @test length(t3) == length(uniq_nums)
+    end
+
+    @testset "in" begin
+        t4 = SplayTree{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 = SplayTree()
+        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 = SplayTree{Int}()
+        for i in 1:10
+            push!(t6, i)
+        end
+        for i in 1:10
+            @test t6[i] == i
+        end
+        @test_throws KeyError getindex(t6, 0)
+        @test_throws KeyError getindex(t6, 11)
+    end
+
+    @testset "key conversion in push!" begin
+        t7 = SplayTree{Int}()
+        push!(t7, Int8(1))
+        @test length(t7) == 1
+        @test haskey(t7, 1)
+    end
+
+    @testset "maximum_node" begin
+        t8 = SplayTree()
+        @test maximum_node(t8.root) == nothing
+        for i in 1:32
+            push!(t8, i)
+        end
+        m1 = maximum_node(t8.root)
+        @test m1.data == 32
+        node = t8.root
+        while node.rightChild != nothing
+            m = maximum_node(node.rightChild)
+            @test m == m1
+            node = node.rightChild
+        end
+    end
+end
\ No newline at end of file