Skip to content

Instantly share code, notes, and snippets.

@PM2Ring
Created January 2, 2018 20:51
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save PM2Ring/08528cfe301ba4fa4f77cd0e4f4511ec to your computer and use it in GitHub Desktop.
Save PM2Ring/08528cfe301ba4fa4f77cd0e4f4511ec to your computer and use it in GitHub Desktop.
A Python version of the Fenwick tree demo code from the Wikipedia article.
#! /usr/bin/env python3
''' Demo of the Fenwick tree, aka Binary Indexed Tree
See https://en.wikipedia.org/wiki/Fenwick_tree
Converted from C to Python by PM 2Ring
'''
def LSB(i):
''' Return the value of the Least Significant Bit position in i '''
return i & -i
def fen_sum(tree, i):
''' Return the sum of the first i elements, 0 through i-1 '''
total = 0
while i:
total += tree[i-1]
i -= LSB(i)
return total
def fen_add(tree, delta, i):
''' Add delta to element i (and thus to fen_sum(tree, j) for all j > i) '''
size = len(tree)
while i < size:
tree[i] += delta
i += LSB(i+1)
def fen_range(tree, i, j):
''' Return the sum of the elements [i..j-1]
Equivalent to fen_sum(tree, j) - fen_sum(tree, i)
'''
total = 0
while j > i:
total += tree[j-1]
j -= LSB(j)
while i > j:
total -= tree[i-1]
i -= LSB(i)
return total
def fen_get(tree, i):
''' Return the value of the element at index i '''
return fen_range(tree, i, i + 1)
def fen_set(tree, value, i):
''' Set the value of the element at index i '''
fen_add(tree, value - fen_get(tree, i), i)
# It's possible to initialize a Fenwick tree using fen_add or
# fen_set (you can even do the latter in-place if you go backward
# through the array), but for bulk initialization, this is faster.
def fen_init(a):
''' Convert a to a Fenwick tree '''
size = len(a)
for i in range(size):
j = i + LSB(i+1)
if j < size:
a[j] += a[i]
# Test
a = [i * 10**i for i in range(1, 10)]
print('Data :', a)
fen_init(a)
print('Tree :', a)
a_sums = [fen_sum(a, i) for i in range(1, len(a) + 1)]
print('Sums :', a_sums)
items = [fen_get(a, i) for i in range(len(a))]
print('Items :', items)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment