Last active
August 29, 2015 14:01
-
-
Save yurivish/aff46c190c1ac538c46f to your computer and use it in GitHub Desktop.
A Julia implementation of a treap (v6; in progress)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Base: isempty, start, next, done, length, getindex, minimum, maximum, search, show | |
export Treap, TreapNode | |
type TreapNode{K} | |
priority::Float64 | |
size::Int | |
key::K | |
left::TreapNode{K} | |
right::TreapNode{K} | |
TreapNode(key::K, priority::Float64, left::TreapNode{K}, right::TreapNode{K}) = | |
new(priority, left.size + right.size + 1, key, left, right) | |
TreapNode(key::K, priority::Float64) = | |
new(priority, 1, key, TreapNode{K}(), TreapNode{K}()) | |
TreapNode() = new(Inf, 0) | |
end | |
isempty(t::TreapNode) = t.size == 0 | |
show(io::IO, t::TreapNode) = show(io, "key:$(t.key) size:$(t.size) priority:$(t.priority)") | |
start(t::TreapNode) = 1 | |
next(t::TreapNode, state) = (t[state], state + 1) | |
done(t::TreapNode, state) = state > length(t) | |
length(t::TreapNode) = t.size | |
key(t::TreapNode) = t.key | |
left(t::TreapNode) = t.left | |
right(t::TreapNode) = t.right | |
function getindex{K}(t::TreapNode{K}, index::Int) | |
1 <= index <= t.size || throw(KeyError(index)) | |
while t.left.size != index - 1 | |
if index <= t.left.size | |
t = t.left | |
else | |
index = index - t.left.size - 1 | |
t = t.right | |
end | |
end | |
t.key | |
end | |
function minimum{K}(t::TreapNode{K}) | |
isempty(t) && error("An empty treap has no minimum.") | |
while !isempty(t.left) t = t.left end | |
t.key | |
end | |
function maximum{K}(t::TreapNode{K}) | |
isempty(t) && error("An empty treap has no maximum.") | |
while !isempty(t.right) t = t.right end | |
t.key | |
end | |
function add!{K}(t::TreapNode{K}, key::K) | |
isempty(t) && return TreapNode{K}(key, rand()) | |
t.size += 1 | |
if key < t.key | |
t.left = add!(t.left, key) | |
t.left.priority < t.priority ? rotate_right!(t) : t | |
else | |
@assert t.key < key "A treap may not contain duplicate keys: $key, $(t.key)" | |
t.right = add!(t.right, key) | |
t.right.priority < t.priority ? rotate_left!(t) : t | |
end | |
end | |
function merge!{K}(left::TreapNode{K}, right::TreapNode{K}) | |
isempty(left) && return right | |
isempty(right) && return left | |
if left.priority < right.priority | |
result = left | |
result.size += right.size | |
result.right = merge!(left.right, right) | |
else | |
result = right | |
result.size += left.size | |
result.left = merge!(left, result.left) | |
end | |
result | |
end | |
function remove!{K}(t::TreapNode{K}, key::K) | |
isempty(t) && throw(KeyError(key)) | |
t.size -= 1 | |
if key == t.key | |
merge!(t.left, t.right) | |
elseif key < t.key | |
t.left = remove!(t.left, key) | |
t.left.priority < t.priority ? rotate_right!(t) : t | |
else | |
t.right = remove!(t.right, key) | |
t.right.priority < t.priority ? rotate_left!(t) : t | |
end | |
end | |
function rotate_right!{K}(root::TreapNode{K}) | |
@assert !isempty(root) | |
newroot = root.left | |
# Rotate | |
root.left = newroot.right | |
newroot.right = root | |
# Update sizes | |
root.size -= newroot.left.size + 1 | |
newroot.size += root.right.size + 1 | |
newroot | |
end | |
function rotate_left!{K}(root::TreapNode{K}) | |
@assert !isempty(root) | |
newroot = root.right | |
# Rotate | |
root.right = newroot.left | |
newroot.left = root | |
# Update sizes | |
root.size -= newroot.right.size + 1 | |
newroot.size += root.left.size + 1 | |
newroot | |
end | |
type Treap{K} | |
root::TreapNode{K} | |
Treap() = new(TreapNode{K}()) | |
end | |
add!{K}(t::Treap{K}, key::K) = t.root = add!(t.root, key) | |
remove!{K}(t::Treap{K}, key::K) = t.root = remove!(t.root, key) | |
length(t::Treap) = length(t.root) | |
getindex(t::Treap, n::Int) = getindex(t.root, n) | |
start(t::Treap) = start(t.root) | |
next(t::Treap, state) = next(t.root, state) | |
done(t::Treap, state) = done(t.root, state) | |
minimum(t::Treap) = minimum(t.root) | |
maximum(t::Treap) = maximum(t.root) | |
key(t::Treap) = key(t.root) | |
left(t::Treap) = left(t.root) | |
right(t::Treap) = right(t.root) | |
root(t::Treap) = t.root | |
function test_treap() | |
n = 10000 | |
a = shuffle([i for i in 1:n]) | |
t = Treap{Int}() | |
for i in 1:n | |
add!(t, a[i]) | |
end | |
sort!(a) | |
@assert !isempty(t) | |
@assert length(t) == length(a) "$(length(t)) != $(length(a))" | |
@assert maximum(t) == maximum(a) | |
@assert minimum(t) == minimum(a) | |
# Check that all elements have been added to | |
# the treap and are in sorted order. | |
for (i, v) in enumerate(a) | |
@assert t[i] == v | |
end | |
for v in a | |
remove!(t, v) | |
end | |
@assert length(t) == 0 | |
@assert isempty(t) | |
println("Treap: Test succeeded.") | |
end | |
function benchmark_treap(n) | |
gc_disable() | |
t = Treap{Int}() | |
a = shuffle([i for i in 1:n]) | |
println("Timing $n insert operations.") | |
@time for i in 1:n | |
t = add!(t, a[i]) | |
end | |
@assert length(t) == n | |
println("Timing $n random access operations.") | |
@time for i in 1:n | |
t[rand(1:n)] | |
end | |
println("Timing $n remove operations.") | |
@time for i in 1:n | |
remove!(t, a[i]) | |
end | |
@assert isempty(t) | |
gc_enable() | |
end | |
# test_treap() | |
# benchmark_treap(100000) | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment