data_structures.binary_tree.fenwick_tree

Classes

FenwickTree

Fenwick Tree

Module Contents

class data_structures.binary_tree.fenwick_tree.FenwickTree(arr: list[int] | None = None, size: int | None = None)

Fenwick Tree

More info: https://en.wikipedia.org/wiki/Fenwick_tree

add(index: int, value: int) → None

Add a value to index in O(lg N)

Parameters:

index (int): index to add value to value (int): value to add to index

Returns:

None

>>> f = FenwickTree([1, 2, 3, 4, 5])
>>> f.add(0, 1)
>>> f.add(1, 2)
>>> f.add(2, 3)
>>> f.add(3, 4)
>>> f.add(4, 5)
>>> f.get_array()
[2, 4, 6, 8, 10]

get(index: int) → int

Get value at index in O(lg N)

Parameters:

index (int): index to get the value

Returns:

int: Value of element at index

>>> a = [i for i in range(128)]
>>> f = FenwickTree(a)
>>> res = True
>>> for i in range(len(a)):
...     res = res and f.get(i) == a[i]
>>> res
True

get_array() → list[int]

Get the Normal Array of the Fenwick tree in O(N)

Returns:

list: Normal Array of the Fenwick tree

>>> a = [i for i in range(128)]
>>> f = FenwickTree(a)
>>> f.get_array() == a
True

init(arr: list[int]) → None

Initialize the Fenwick tree with arr in O(N)

Parameters:

arr (list): list of elements to initialize the tree with

Returns:

None

>>> a = [1, 2, 3, 4, 5]
>>> f1 = FenwickTree(a)
>>> f2 = FenwickTree(size=len(a))
>>> for index, value in enumerate(a):
...     f2.add(index, value)
>>> f1.tree == f2.tree
True

static next_(index: int) → int

prefix(right: int) → int

Prefix sum of all elements in [0, right) in O(lg N)

Parameters:

right (int): right bound of the query (exclusive)

Returns:

int: sum of all elements in [0, right)

>>> a = [i for i in range(128)]
>>> f = FenwickTree(a)
>>> res = True
>>> for i in range(len(a)):
...     res = res and f.prefix(i) == sum(a[:i])
>>> res
True

static prev(index: int) → int

query(left: int, right: int) → int

Query the sum of all elements in [left, right) in O(lg N)

Parameters:

left (int): left bound of the query (inclusive) right (int): right bound of the query (exclusive)

Returns:

int: sum of all elements in [left, right)

>>> a = [i for i in range(128)]
>>> f = FenwickTree(a)
>>> res = True
>>> for i in range(len(a)):
...     for j in range(i + 1, len(a)):
...         res = res and f.query(i, j) == sum(a[i:j])
>>> res
True

rank_query(value: int) → int

Find the largest index with prefix(i) <= value in O(lg N) NOTE: Requires that all values are non-negative!

Parameters:

value (int): value to find the largest index of

Returns:

-1: if value is smaller than all elements in prefix sum int: largest index with prefix(i) <= value

>>> f = FenwickTree([1, 2, 0, 3, 0, 5])
>>> f.rank_query(0)
-1
>>> f.rank_query(2)
0
>>> f.rank_query(1)
0
>>> f.rank_query(3)
2
>>> f.rank_query(5)
2
>>> f.rank_query(6)
4
>>> f.rank_query(11)
5

update(index: int, value: int) → None

Set the value of index in O(lg N)

Parameters:

index (int): index to set value to value (int): value to set in index

Returns:

None

>>> f = FenwickTree([5, 4, 3, 2, 1])
>>> f.update(0, 1)
>>> f.update(1, 2)
>>> f.update(2, 3)
>>> f.update(3, 4)
>>> f.update(4, 5)
>>> f.get_array()
[1, 2, 3, 4, 5]