fenwick tree more info https en wikipedia orgwikifenwicktree constructor for the fenwick tree parameters arr list list of elements to initialize the tree with optional size int size of the fenwick tree if arr is none initialize the fenwick tree with arr in on parameters arr list list of elements to initialize the tree with returns none a 1 2 3 4 5 f1 fenwicktreea f2 fenwicktreesizelena for index value in enumeratea f2 addindex value f1 tree f2 tree true get the normal array of the fenwick tree in on returns list normal array of the fenwick tree a i for i in range128 f fenwicktreea f getarray a true add a value to index in olg n parameters index int index to add value to value int value to add to index returns none f fenwicktree1 2 3 4 5 f add0 1 f add1 2 f add2 3 f add3 4 f add4 5 f getarray 2 4 6 8 10 set the value of index in olg n parameters index int index to set value to value int value to set in index returns none f fenwicktree5 4 3 2 1 f update0 1 f update1 2 f update2 3 f update3 4 f update4 5 f getarray 1 2 3 4 5 prefix sum of all elements in 0 right in olg n parameters right int right bound of the query exclusive returns int sum of all elements in 0 right a i for i in range128 f fenwicktreea res true for i in rangelena res res and f prefixi suma i res true query the sum of all elements in left right in olg 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 range128 f fenwicktreea res true for i in rangelena for j in rangei 1 lena res res and f queryi j sumai j res true get value at index in olg n parameters index int index to get the value returns int value of element at index a i for i in range128 f fenwicktreea res true for i in rangelena res res and f geti ai res true find the largest index with prefixi value in olg n note requires that all values are nonnegative 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 prefixi value f fenwicktree1 2 0 3 0 5 f rankquery0 1 f rankquery2 0 f rankquery1 0 f rankquery3 2 f rankquery5 2 f rankquery6 4 f rankquery11 5 fenwick tree more info https en wikipedia org wiki fenwick_tree constructor for the fenwick tree parameters arr list list of elements to initialize the tree with optional size int size of the fenwick tree if arr is 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 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 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 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 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 make right inclusive 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 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 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 largest power of 2 size
from copy import deepcopy


class FenwickTree:
    

    def __init__(self, arr: list[int] | None = None, size: int | None = None) -> None:
        

        if arr is None and size is not None:
            self.size = size
            self.tree = [0] * size
        elif arr is not None:
            self.init(arr)
        else:
            raise ValueError("Either arr or size must be specified")

    def init(self, arr: list[int]) -> None:
        
        self.size = len(arr)
        self.tree = deepcopy(arr)
        for i in range(1, self.size):
            j = self.next_(i)
            if j < self.size:
                self.tree[j] += self.tree[i]

    def get_array(self) -> list[int]:
        
        arr = self.tree[:]
        for i in range(self.size - 1, 0, -1):
            j = self.next_(i)
            if j < self.size:
                arr[j] -= arr[i]
        return arr

    @staticmethod
    def next_(index: int) -> int:
        return index + (index & (-index))

    @staticmethod
    def prev(index: int) -> int:
        return index - (index & (-index))

    def add(self, index: int, value: int) -> None:
        
        if index == 0:
            self.tree[0] += value
            return
        while index < self.size:
            self.tree[index] += value
            index = self.next_(index)

    def update(self, index: int, value: int) -> None:
        
        self.add(index, value - self.get(index))

    def prefix(self, right: int) -> int:
        
        if right == 0:
            return 0
        result = self.tree[0]
        right -= 1  
        while right > 0:
            result += self.tree[right]
            right = self.prev(right)
        return result

    def query(self, left: int, right: int) -> int:
        
        return self.prefix(right) - self.prefix(left)

    def get(self, index: int) -> int:
        
        return self.query(index, index + 1)

    def rank_query(self, value: int) -> int:
        
        value -= self.tree[0]
        if value < 0:
            return -1

        j = 1  
        while j * 2 < self.size:
            j *= 2

        i = 0

        while j > 0:
            if i + j < self.size and self.tree[i + j] <= value:
                value -= self.tree[i + j]
                i += j
            j //= 2
        return i


if __name__ == "__main__":
    import doctest

    doctest.testmod()