Generic Fenwick Tree

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) }

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

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.

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))
additiveTree.accumulated(in: 1...3))
additiveTree.accumulated(in: 2...))
additiveTree.accumulated(in: ..<4))
additiveTree.accumulated(in: ...3))

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[0] 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[0]
    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.