Skip to content

Commit

Permalink
Add Fenwick Tree (#485)
Browse files Browse the repository at this point in the history
Add Fenwick Tree

* Methods for manipulating/accessing the tree: `inc!`, `dec!`, `incdec!`, and `prefixsum`
  • Loading branch information
eulerkochy authored and kmsquire committed Apr 8, 2019
1 parent 69c867f commit e0cc5d0
Show file tree
Hide file tree
Showing 4 changed files with 139 additions and 2 deletions.
6 changes: 5 additions & 1 deletion src/DataStructures.jl
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ module DataStructures
ReverseOrdering, Reverse, Lt,
isless, union, intersect, symdiff, setdiff, issubset,
searchsortedfirst, searchsortedlast, in,
eachindex, keytype, valtype, minimum, maximum
eachindex, keytype, valtype, minimum, maximum, size

using OrderedCollections
import OrderedCollections: filter, filter!, isordered
Expand All @@ -28,6 +28,8 @@ module DataStructures
export classified_lists, classified_sets, classified_counters

export IntDisjointSets, DisjointSets, num_groups, find_root, in_same_set, root_union!

export FenwickTree, length, inc!, dec!, incdec!, prefixsum

export AbstractHeap, compare, extract_all!
export BinaryHeap, BinaryMinHeap, BinaryMaxHeap, nlargest, nsmallest
Expand Down Expand Up @@ -73,6 +75,8 @@ module DataStructures

include("int_set.jl")

include("fenwick.jl")

include("list.jl")
include("mutable_list.jl")
include("balanced_tree.jl")
Expand Down
95 changes: 95 additions & 0 deletions src/fenwick.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
"""
FenwickTree{T}(n)
Constructs a [`FenwickTree`](https://en.wikipedia.org/wiki/Fenwick_tree) of length `n`.
"""
struct FenwickTree{T}
bi_tree::Vector{T} #bi_tree is shorthand for Binary Indexed Tree, an alternative name for Fenwick Tree
n::Int
end

FenwickTree{T}(n::Integer) where T = FenwickTree{T}(zeros(T, n), n)

"""
FenwickTree(arr::AbstractArray)
Constructs a [`FenwickTree`](https://en.wikipedia.org/wiki/Fenwick_tree) from an array of `counts`
"""
function FenwickTree(a::AbstractVector{U}) where U
n = length(a)
tree = FenwickTree{U}(n)
@inbounds for i = 1:n
inc!(tree, i, a[i])
end
tree
end

length(ft::FenwickTree) = ft.n

"""
inc!(ft::FenwickTree{T}, ind, val)
Increases the value of the [`FenwickTree`] by `val` from the index `ind` upto the length of the Fenwick Tree.
"""
function inc!(ft::FenwickTree{T}, ind::Integer, val = 1) where T
val0 = convert(T, val)
i = ind
n = ft.n
@boundscheck 1 <= i <= n || throw(ArgumentError("$i should be in between 1 and $n"))
@inbounds while i <= n
ft.bi_tree[i] += val0
i += i&(-i)
end
end

"""
dec!(ft::FenwickTree, ind, val)
Decreases the value of the [`FenwickTree`] by `val` from the index `ind` upto the length of the Fenwick Tree.
"""
dec!(ft::FenwickTree, ind::Integer, val = 1 ) = inc!(ft, ind, -val)

"""
incdec!(ft::FenwickTree{T}, left, right, val)
Increases the value of the [`FenwickTree`] by `val` from the indices from `left` and decreases it from the `right`.
"""
function incdec!(ft::FenwickTree{T}, left::Integer, right::Integer, val = one(T)) where T
val0 = convert(T, val)
inc!(ft, left, val0)
dec!(ft, right, val0)
end

"""
prefixsum(ft::FenwickTree{T}, ind)
Return the cumulative sum from index 1 upto `ind` of the [`FenwickTree`](@ref)
# Examples
```
julia> f = FenwickTree{Int}(6)
julia> inc!(f, 2, 5)
julia> prefixsum(f, 1)
0
julia> prefixsum(f, 3)
5
```
"""
function prefixsum(ft::FenwickTree{T}, ind::Integer) where T
sum = zero(T)
i = ind
n = ft.n
@boundscheck 1 <= i <= n || throw(ArgumentError("$i should be in between 1 and $n"))
@inbounds while i > 0
sum += ft.bi_tree[i]
i -= i&(-i)
end
sum
end

getindex(ft::FenwickTree{T}, ind::Integer) where T = prefixsum(ft, ind)
3 changes: 2 additions & 1 deletion test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@ tests = ["int_set",
"multi_dict",
"circular_buffer",
"sorting",
"priority_queue"
"priority_queue",
"fenwick"
]

if length(ARGS) > 0
Expand Down
37 changes: 37 additions & 0 deletions test/test_fenwick.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
@testset "Fenwick Tree" begin
@testset "initialisation" begin
f1 = FenwickTree{Float64}(5)
@test f1.bi_tree == zeros(5)
@test length(f1) == 5

arr = [1.2, 8.7, 7.2, 3.5]
f3 = FenwickTree(arr)
@test f3[1] == 1.2
@test length(f3) == length(arr)
end

@testset "Point update and Point queries" begin
f1 = FenwickTree{Int}(10)
inc!(f1, 10, 5)
@test prefixsum(f1, 10) == 5
@test prefixsum(f1, 9) == 0
@test prefixsum(f1, 1) == 0
inc!(f1, 5, 7)
@test prefixsum(f1, 8) == 7
@test prefixsum(f1, 10) == 12
@test prefixsum(f1, 5) == 7
@test prefixsum(f1, 4) == 0

dec!(f1, 7, 2)
@test prefixsum(f1, 8) == 5
@test prefixsum(f1, 6) == 7

incdec!(f1, 2, 6, 3)
@test prefixsum(f1, 3) == 3
@test prefixsum(f1, 7) == 5

@test_throws ArgumentError inc!(f1, 11)
@test_throws ArgumentError inc!(f1, 0)
end

end

0 comments on commit e0cc5d0

Please sign in to comment.