-
Notifications
You must be signed in to change notification settings - Fork 245
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add Fenwick Tree * Methods for manipulating/accessing the tree: `inc!`, `dec!`, `incdec!`, and `prefixsum`
- Loading branch information
1 parent
69c867f
commit e0cc5d0
Showing
4 changed files
with
139 additions
and
2 deletions.
There are no files selected for viewing
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
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
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) |
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
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
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 |