Fix maximum Fenwick tree range values after updates

data_structures/binary_tree/maximum_fenwick_tree.py

@@ -34,6 +34,26 @@ class MaxFenwickTree:
     >>> ft.update(0, 1000)
     >>> ft.query(0, 1)
     1000
+
+    Updating a smaller sibling preserves the maximum, including after decreases.
+    >>> ft = MaxFenwickTree(8)
+    >>> ft.update(4, 20)
+    >>> ft.update(5, 1)
+    >>> ft.query(0, 6)
+    20
+    >>> ft.update(4, 0)
+    >>> ft.query(0, 6)
+    1
+    >>> ft.update(5, 0)
+    >>> ft.query(0, 6)
+    0
+
+    Updating a child must also preserve the ancestor's own array value.
+    >>> ft = MaxFenwickTree(8)
+    >>> ft.update(5, 100)
+    >>> ft.update(4, 0)
+    >>> ft.query(0, 6)
+    100
     """
 
     def __init__(self, size: int) -> None:
@@ -75,13 +95,32 @@ def update(self, index: int, value: int) -> None:
         Returns:
             None
         """
+        old_value = self.arr[index]
+        if value == old_value:
+            return
         self.arr[index] = value
         while index < self.size:
-            current_left_border = self.get_prev(index) + 1
-            if current_left_border == index:
+            old_maximum = self.tree[index]
+            if value > old_maximum:
                 self.tree[index] = value
+            elif old_value == old_maximum:
+                current_left_border = self.get_prev(index) + 1
+                maximum = self.arr[index]
+                if current_left_border != index:
+                    maximum = max(0, maximum)
+                child = index - 1
+                # These disjoint child buckets cover the rest of this bucket.
+                while child >= current_left_border:
+                    maximum = max(maximum, self.tree[child])
+                    if maximum == old_maximum:
+                        break
+                    child = self.get_prev(child)
+                self.tree[index] = maximum
+                if maximum == old_maximum:
+                    break
             else:
-                self.tree[index] = max(value, current_left_border, index)
+                # An unchanged bucket maximum leaves all its ancestors unchanged.
+                break
             index = self.get_next(index)
 
     def query(self, left: int, right: int) -> int:
@@ -110,5 +149,41 @@ def query(self, left: int, right: int) -> int:
 
 if __name__ == "__main__":
     import doctest
+    import sys
+    from timeit import repeat
 
     doctest.testmod()
+
+    # Run with --benchmark on each revision to compare the same 2,000-item workload.
+    if "--benchmark" in sys.argv:
+        size = 2000
+        values = [(index * 97) % size for index in range(size)]
+        updates = [(index, (index * 37) % size) for index in range(size)]
+        updates += [(index, values[index]) for index in reversed(range(size))]
+        queries = [(left, size) for left in range(size)]
+
+        # Build correct query buckets outside the timer, even on the buggy revision.
+        query_setup = """
+tree = MaxFenwickTree(size)
+tree.arr = values[:]
+tree.tree = [
+    max(values[tree.get_prev(index) + 1 : index + 1])
+    for index in range(size)
+]
+"""
+        for operation, statement, setup in (
+            (
+                "4000 updates",
+                "for index, value in updates: tree.update(index, value)",
+                "tree = MaxFenwickTree(size)",
+            ),
+            (
+                "2000 queries",
+                "for left, right in queries: tree.query(left, right)",
+                query_setup,
+            ),
+        ):
+            timings = repeat(
+                statement, setup=setup, repeat=5, number=1, globals=globals()
+            )
+            print(f"{size} items, {operation}: {min(timings):.6f} seconds (best of 5)")