Python Fenwick Tree

Question

I am trying to solve this question: https://open.kattis.com/problems/fenwick

A Fenwick Tree (also known as a Binary Indexed Tree) is a data structure on an array which enables fast (𝑂(log𝑛)) updates and prefix sum queries on the underlying data.

For this problem, implement a Fenwick Tree to support operations of two types: (a) increment an element in the array or (b) query the prefix sum of a portion of the array.

Input The first line of input contains two integers 𝑁, 𝑄, where 1≤𝑁≤5000000 is the length of the array and 0≤𝑄≤5000000 is the number of operations. Then follow 𝑄 lines giving the operations. There are two types of operations:

“+ 𝑖 𝛿” indicates that 𝑎[𝑖] is incremented by 𝛿, where 0≤𝑖<𝑁 and −10^9≤𝛿≤10^9 (both are integers)

“? 𝑖” is a query for the value of 𝑎[0]+𝑎[1]+…+𝑎[𝑖−1], where 0≤𝑖≤𝑁 (for 𝑖=0 this is interpreted as an empty sum)

You should assume that every array entry is initially 0.

My implementation:

class SegmentTree:
    def __init__(self, n):
        self.nums = [0] * n
        self.tree = {}
        self._build()

    def _build(self):
        def build(l, r, index):
            if l == r:
                self.tree[index] = self.nums[l]
            else:
                mid = (l+r)//2
                build(l, mid, index*2)
                build(mid+1, r, index*2+1)
                self.tree[index] = self.tree[index*2] + self.tree[index*2+1]

        build(0, len(self.nums)-1, 1)

    def get_sum(self, i, j):
        def get_sum(left_boundary, right_boundary, index, i, j):
            if j < i:
                return 0

            if left_boundary == right_boundary:
                return self.tree[index]
            if left_boundary == i and right_boundary == j:
                return self.tree[index]
            else:
                mid = (left_boundary + right_boundary)//2
                left_sum = get_sum(left_boundary, mid, index*2, i, min(j, mid))
                right_sum = get_sum(mid+1, right_boundary,
                                    index*2+1, max(i, mid+1), j)

                return left_sum + right_sum

        return get_sum(0, len(self.nums)-1, 1, i, j)

    def update(self, pos, num):
        def update(l, r, index):
            if l == r and l == pos:
                self.tree[index] = num
            else:
                mid = (l+r)//2
                if pos <= mid:
                    update(l, mid, index*2)
                else:
                    update(mid+1, r, index*2+1)
                self.tree[index] = self.tree[index*2] + self.tree[index*2+1]
        update(0, len(self.nums)-1, 1)


def main():

    n, q = map(int, input().split())
    nums = [0] * n
    s = SegmentTree(n)

    def update(ops):
        plus, i, inc = ops.split()
        i = int(i)
        inc = int(inc)
        s.nums[i] += inc
        s.update(i, s.nums[i])

    def prefix_sum(ops):
        q, j = ops.split()
        j = int(j)
        ans = s.get_sum(0, j-1)
        print(ans)

    for _ in range(q):
        ops = input()
        if len(ops.split()) == 3:
            update(ops)
        else:
            prefix_sum(ops)


if __name__ == "__main__":
    main()

This sadly times out. How do I optimize it to run faster?

Answer 1

Segment trees and Fenwick trees are usually implemented as implicit data structures. That is, as an array with the tree structure implicitly given by the array indices. Your code instead stores the tree as a Python dictionary which is very inefficient.

The second problem with your segment tree is that you are using recursion. Refer to this sample code for how to implement iterative versions of get_sum and update.

The third problem is that you have implemented a Segment tree and not a Fenwick tree (aka Binary indexed tree). For this problem, they are slightly more efficient than Segment trees.

The fourth problem is, sadly, Python. :) Number crunching code and tight loops just doesn't run as fast in Python as they do in lower-level languages. Plus, in this problem, IO and number parsing is dominating.

Here is a Python solution using a Fenwick tree:

n, q = map(int, input().split())
# One extra element so that we can use 1-based indexing.
nums = [0] * (n + 1)
for _ in range(q):
    parts = input().split()
    if len(parts) == 3:
        _, i, inc = parts[0], int(parts[1]) + 1, int(parts[2])
        while i <= n:
            nums[i] += inc
            i += i & (-i)
    else:
        _, i = parts[0], int(parts[1])
        tot = 0
        while i > 0:
            tot += nums[i]
            i -= i & (-i)
        print(tot)

If you port it to a low-level language it will be very fast.