@@ -16,7 +16,7 @@ class MaxFenwickTree:
1616 20
1717 >>> ft.update(4, 10)
1818 >>> ft.query(2, 5)
19- 10
19+ 20
2020 >>> ft.query(1, 5)
2121 20
2222 >>> ft.update(2, 0)
@@ -26,6 +26,14 @@ class MaxFenwickTree:
2626 >>> ft.update(255, 30)
2727 >>> ft.query(0, 10000)
2828 30
29+ >>> ft = MaxFenwickTree(6)
30+ >>> ft.update(5, 1)
31+ >>> ft.query(5, 6)
32+ 1
33+ >>> ft = MaxFenwickTree(6)
34+ >>> ft.update(0, 1000)
35+ >>> ft.query(0, 1)
36+ 1000
2937 """
3038
3139 def __init__ (self , size : int ) -> None :
@@ -47,14 +55,14 @@ def get_next(index: int) -> int:
4755 """
4856 Get next index in O(1)
4957 """
50- return index + (index & - index )
58+ return index | (index + 1 )
5159
5260 @staticmethod
5361 def get_prev (index : int ) -> int :
5462 """
5563 Get previous index in O(1)
5664 """
57- return index - (index & - index )
65+ return ( index & (index + 1 )) - 1
5866
5967 def update (self , index : int , value : int ) -> None :
6068 """
@@ -69,7 +77,11 @@ def update(self, index: int, value: int) -> None:
6977 """
7078 self .arr [index ] = value
7179 while index < self .size :
72- self .tree [index ] = max (value , self .query (self .get_prev (index ), index ))
80+ current_left_border = self .get_prev (index ) + 1
81+ if current_left_border == index :
82+ self .tree [index ] = value
83+ else :
84+ self .tree [index ] = max (value , current_left_border , index )
7385 index = self .get_next (index )
7486
7587 def query (self , left : int , right : int ) -> int :
@@ -85,9 +97,9 @@ def query(self, left: int, right: int) -> int:
8597 """
8698 right -= 1 # Because of right is exclusive
8799 result = 0
88- while left < right :
100+ while left <= right :
89101 current_left = self .get_prev (right )
90- if left < current_left :
102+ if left <= current_left :
91103 result = max (result , self .tree [right ])
92104 right = current_left
93105 else :
0 commit comments