Python/data_structures/binary_tree/maximum_fenwick_tree.py

115 lines
2.7 KiB
Python
Raw Normal View History

class MaxFenwickTree:
"""
Maximum Fenwick Tree
More info: https://cp-algorithms.com/data_structures/fenwick.html
---------
>>> ft = MaxFenwickTree(5)
>>> ft.query(0, 5)
0
>>> ft.update(4, 100)
>>> ft.query(0, 5)
100
>>> ft.update(4, 0)
>>> ft.update(2, 20)
>>> ft.query(0, 5)
20
>>> ft.update(4, 10)
>>> ft.query(2, 5)
2022-09-14 08:24:55 +00:00
20
>>> ft.query(1, 5)
20
>>> ft.update(2, 0)
>>> ft.query(0, 5)
10
>>> ft = MaxFenwickTree(10000)
>>> ft.update(255, 30)
>>> ft.query(0, 10000)
30
2022-09-14 08:24:55 +00:00
>>> ft = MaxFenwickTree(6)
>>> ft.update(5, 1)
>>> ft.query(5, 6)
1
>>> ft = MaxFenwickTree(6)
>>> ft.update(0, 1000)
>>> ft.query(0, 1)
1000
"""
def __init__(self, size: int) -> None:
"""
Create empty Maximum Fenwick Tree with specified size
Parameters:
size: size of Array
Returns:
None
"""
self.size = size
self.arr = [0] * size
self.tree = [0] * size
@staticmethod
def get_next(index: int) -> int:
"""
Get next index in O(1)
"""
2022-09-14 08:24:55 +00:00
return index | (index + 1)
@staticmethod
def get_prev(index: int) -> int:
"""
Get previous index in O(1)
"""
2022-09-14 08:24:55 +00:00
return (index & (index + 1)) - 1
def update(self, index: int, value: int) -> None:
"""
Set index to value in O(lg^2 N)
Parameters:
index: index to update
value: value to set
Returns:
None
"""
self.arr[index] = value
while index < self.size:
2022-09-14 08:24:55 +00:00
current_left_border = self.get_prev(index) + 1
if current_left_border == index:
self.tree[index] = value
else:
self.tree[index] = max(value, current_left_border, index)
index = self.get_next(index)
def query(self, left: int, right: int) -> int:
"""
Answer the query of maximum range [l, r) in O(lg^2 N)
Parameters:
left: left index of query range (inclusive)
right: right index of query range (exclusive)
Returns:
Maximum value of range [left, right)
"""
right -= 1 # Because of right is exclusive
result = 0
2022-09-14 08:24:55 +00:00
while left <= right:
current_left = self.get_prev(right)
2022-09-14 08:24:55 +00:00
if left <= current_left:
result = max(result, self.tree[right])
right = current_left
else:
result = max(result, self.arr[right])
right -= 1
return result
if __name__ == "__main__":
import doctest
doctest.testmod()