# Generic Fenwick Tree

Posted on

Problem

A Fenwick Tree, or a Binary Indexed tree, is an interesting data structure that can efficiently update its elements. You can read more about it in this paper.

Here is a generic implementation of a Fenwick Tree, which could take any two reversible operations:

``````class FenwickTree<T> {

private let count   : Int
private let neutral : T
private let forward : (T,T) -> T
private let reverse : (T,T) -> T

private var data    : [T]
private var tree    : [T]

init(
count: Int,
neutralElement: T,
forward: @escaping (T,T) -> T,
reverse: @escaping (T,T) -> T
) {
self.count = count
self.neutral = neutralElement
self.forward = forward
self.reverse = reverse
self.data = Array(repeating: neutralElement, count: count)
self.tree = Array(repeating: neutralElement, count: count + 1)
}

func update(index: Int, with newValue: T) {
let oldValue = data[index];
let delta = reverse(newValue, oldValue)
data[index] = newValue
var treeIndex = index + 1
while treeIndex <= count {
tree[treeIndex] = forward(tree[treeIndex], delta)
treeIndex += treeIndex & -treeIndex
}
}

func accumulated(at index: Int) -> T {
var sum = neutral
var treeIndex = index + 1
while 0 < treeIndex {
sum = forward(tree[treeIndex], sum)
treeIndex -= treeIndex & -treeIndex
}
return sum
}

func accumulated(in range: Range<Int>) -> T {
let low = range.lowerBound, high = range.upperBound - 1
let cumulatedLow = low == 0 ? neutral : accumulated(at: low - 1)
let cumulatedHigh = accumulated(at: high)
return low == high ? data[low] : reverse(cumulatedHigh,cumulatedLow)
}
}
``````

Here are some example usecases:

``````let n = 10

let additiveTree = FenwickTree(count: n, neutralElement: 0, forward: +, reverse: -)

for i in 0..<n { additiveTree.update(index: i, with: i + 1) }

let multiplicativeTree = FenwickTree(count: n, neutralElement: 1, forward: *, reverse: /)

for i in 0..<n { multiplicativeTree.update(index: i, with: i + 1) }

multiplicativeTree.accumulated(at: 5)
multiplicativeTree.accumulated(at: 9)
multiplicativeTree.accumulated(in: 8..<10)
``````

I am mainly concerned about the naming of functions and variables:

• `neutral`: This is the neutral element for `forward` and `reverse`
• `forward` and `reverse` should be transitive verbs since they have side effects. Preferably, the names should have the same number of letters.
• `accumulated(at:)` is an unusual name, more familiar would be `foldLeft` from Scala, or `foldl` from Haskell.
• `accumulated(at:)` and `accumulated(in:)` have too similar names.

I can see some possible improvements:

• Performance: `accumulated(in:)` could have better time complexity at the expense of some auxiliary space;
• Enforcing the reversibility of `forward` and `reverse` in code;

Other areas of improvement are welcome.

Solution

### Choosing the generic type (and parameters)

The Fenwick tree stores cumulative values of its elements, and is based on the ability to compute the sum and difference of those elements. You made a very general implementation, where “sum” and “difference” are provided as closure parameters. But it cannot be ensured that these are proper inverse operations of each other, and a slight variation of your `multiplicativeTree` example demonstrates the problem:

``````let multiplicativeTree = FenwickTree(count: 4, neutralElement: 1,
forward: *, reverse: /)

multiplicativeTree.update(index: 0, with: 5)
print(multiplicativeTree.accumulated(at: 0))    // 5   OK

multiplicativeTree.update(index: 0, with: 10)
print(multiplicativeTree.accumulated(at: 0))    // 10  still OK

multiplicativeTree.update(index: 0, with: 13)
print(multiplicativeTree.accumulated(at: 0))    // 10  what?

multiplicativeTree.update(index: 0, with: 11)
print(multiplicativeTree.accumulated(at: 0))    // 0   what???
``````

I would require the elements to conform to the `AdditiveArithmetic` protocol instead. That protocol

• requires addition and subtraction, making the `forward` and `reverse` parameter obsolete,
• requires a `.zero` element, making the `neutralElement` parameter obsolete, and
• provides semantics: according to SE-0223, it “roughly corresponds to the mathematic notion of an additive group.”

`AdditiveArithmetic` is adopted by all binary integer and floating point types, that probably covers most use-cases. If you really need a data structure where updating a value is done via multiplication then you can define a custom type which implements `AdditiveArithmetic` accordingly.

The declaration and initialization then simplifies to

``````class FenwickTree<T: AdditiveArithmetic>  {

private let count   : Int

private var data    : [T]
private var tree    : [T]

init(count: Int) {
self.count = count
self.data = Array(repeating: .zero, count: count)
self.tree = Array(repeating: .zero, count: count + 1)
}

// ...
}
``````

and sums/differences can just be computed with `+` and `-`, making the code even better readable.

### Accessing the elements

There is a

``````func update(index: Int, with newValue: T) {
``````

to set new value at an index, but no method to get a single value. I would suggest to implement the `subscript` method instead, with both a getter and a setter:

``````subscript(index: Int) -> T {
get {
return data[index]
}
set {
// ... update `data` and `tree`
}
}
``````

And not much more work is needed to make `FenwickTree` a (mutable, random access) `Collection`.

### A wider range of ranges …

The `accumulated(in:)` method becomes much more flexible if it takes a `RangeExpression` parameter instead of a `Range`:

``````func accumulated<R>(in range: R) -> T
where R: RangeExpression, R.Bound == Int
{
let realRange = range.relative(to: 0..<count)
let cumulatedLow = accumulated(at: realRange.lowerBound - 1)
let cumulatedHigh = accumulated(at: realRange.upperBound - 1)
return cumulatedHigh - cumulatedLow
}
``````

Now you can call it with half-open, closed, or partial ranges:

``````additiveTree.accumulated(in: 1..<2))
``````

The last line is equivalent to your

``````additiveTree.accumulated(at: 3))
``````

which means that `accumulated(at:)` is no longer needed as a public method (which solves the problem of the similar names of those methods.)

### Further suggestions

The index operations

`````` treeIndex += treeIndex & -treeIndex
treeIndex -= treeIndex & -treeIndex
``````

look like magic to the first reader of your code. I would move that to helper methods (which the compiler inlines) and document it.

The `tree` element is never accessed. With slight modifications of the index calculations you can save one element in that array, i.e. initialize it as

``````  self.tree = Array(repeating: .zero, count: count)
``````

For example, the accumulation would be

``````func accumulated(at index: Int) -> T {
guard index >= 0 else { return .zero }
var sum = tree
var treeIndex = index
while treeIndex > 0 {
sum += tree[treeIndex]
treeIndex -= treeIndex & -treeIndex
}
return sum
}
``````

Finally, a single element can be retrieved from the `tree` array alone, by computing the difference of two tree nodes (this is also described in the Fenwick article). Therefore, at the cost of slightly more computation, you can get rid of the `self.data` array and save memory.