From e0cc5d0507648618a0aac9ce382f05c2fc732525 Mon Sep 17 00:00:00 2001
From: Koustav Chowdhury <kc99.kol@gmail.com>
Date: Mon, 8 Apr 2019 11:28:38 +0530
Subject: [PATCH] Add Fenwick Tree (#485)

Add Fenwick Tree

* Methods for manipulating/accessing the tree: `inc!`, `dec!`, `incdec!`, and `prefixsum`
---
 src/DataStructures.jl |  6 ++-
 src/fenwick.jl        | 95 +++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl      |  3 +-
 test/test_fenwick.jl  | 37 +++++++++++++++++
 4 files changed, 139 insertions(+), 2 deletions(-)
 create mode 100644 src/fenwick.jl
 create mode 100644 test/test_fenwick.jl

diff --git a/src/DataStructures.jl b/src/DataStructures.jl
index 8117cc1a7..fe2ca5515 100644
--- a/src/DataStructures.jl
+++ b/src/DataStructures.jl
@@ -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
@@ -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
@@ -73,6 +75,8 @@ module DataStructures
 
     include("int_set.jl")
 
+    include("fenwick.jl")
+
     include("list.jl")
     include("mutable_list.jl")
     include("balanced_tree.jl")
diff --git a/src/fenwick.jl b/src/fenwick.jl
new file mode 100644
index 000000000..4ce221941
--- /dev/null
+++ b/src/fenwick.jl
@@ -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)
diff --git a/test/runtests.jl b/test/runtests.jl
index a9ebbf958..df24b9d2d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -26,7 +26,8 @@ tests = ["int_set",
          "multi_dict",
          "circular_buffer",
          "sorting",
-         "priority_queue"
+         "priority_queue", 
+         "fenwick"
         ]
 
 if length(ARGS) > 0
diff --git a/test/test_fenwick.jl b/test/test_fenwick.jl
new file mode 100644
index 000000000..144e5eb81
--- /dev/null
+++ b/test/test_fenwick.jl
@@ -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
\ No newline at end of file