Refactor LazySegmentTree to be easier to modify

pyrival/data_structures/LazySegmentTree.py

@@ -1,83 +1,118 @@
 class LazySegmentTree:
+    """Lazy segment tree with source-editable commutative updates and aggregates."""
+
     def __init__(self, data, default=0, func=max):
         """initialize the lazy segment tree with data"""
         self._default = default
         self._func = func
 
         self._len = len(data)
         self._size = _size = 1 << (self._len - 1).bit_length()
-        self._lazy = [0] * (2 * _size)
+        self._buffer_idx = 2 * _size
 
-        self.data = [default] * (2 * _size)
+        # For custom operations, replace this state and the four methods below.
+        # Multi-field state also requires changing the marked buffer operations.
+        self._lazy = [0] * (2 * _size + 1)
+        self.data = [default] * (2 * _size + 1)
         self.data[_size:_size + self._len] = data
-        for i in reversed(range(_size)):
-            self.data[i] = func(self.data[i + i], self.data[i + i + 1])
+        for i in range(_size - 1, 0, -1):
+            self._merge_data(i + i, i + i + 1, i)
 
     def __len__(self):
         return self._len
 
-    def _push(self, idx):
-        """push query on idx to its children"""
-        # Let the children know of the queries
-        q, self._lazy[idx] = self._lazy[idx], 0
+    def _get_range(self, idx):
+        """return the leaf range represented by idx"""
+        shift = self._size.bit_length() - idx.bit_length()
+        return idx << shift, (idx + 1) << shift
+
+    def _unset_lazy(self, idx):
+        """clear the lazy update at idx"""
+        self._lazy[idx] = 0
+
+    def _apply_to_data(self, update_idx, data_idx):
+        """apply the lazy update at update_idx to data_idx"""
+        self.data[data_idx] += self._lazy[update_idx]
 
-        self._lazy[2 * idx] += q
-        self._lazy[2 * idx + 1] += q
-        self.data[2 * idx] += q
-        self.data[2 * idx + 1] += q
+    def _apply_to_lazy(self, update_idx, lazy_idx):
+        """compose the update at update_idx into lazy_idx"""
+        self._lazy[lazy_idx] += self._lazy[update_idx]
+
+    def _merge_data(self, left, right, target):
+        """merge the data at left and right into target"""
+        self.data[target] = self._func(self.data[left], self.data[right])
 
     def _update(self, idx):
-        """updates the node idx to know of all queries applied to it via its ancestors"""
-        for i in reversed(range(1, idx.bit_length())):
-            self._push(idx >> i)
+        """apply all updates stored above idx"""
+        for i in range(idx.bit_length() - 1, 0, -1):
+            parent = idx >> i
+            left = parent + parent
+            self._apply_to_data(parent, left)
+            self._apply_to_lazy(parent, left)
+            self._apply_to_data(parent, left + 1)
+            self._apply_to_lazy(parent, left + 1)
+            self._unset_lazy(parent)
 
     def _build(self, idx):
-        """make the changes to idx be known to its ancestors"""
+        """make the changes to idx known to its ancestors"""
         idx >>= 1
         while idx:
-            self.data[idx] = self._func(self.data[2 * idx], self.data[2 * idx + 1]) + self._lazy[idx]
+            self._merge_data(idx + idx, idx + idx + 1, idx)
+            self._apply_to_data(idx, idx)
             idx >>= 1
 
-    def add(self, start, stop, value):
-        """lazily add value to [start, stop)"""
+    def apply(self, start, stop, value):
+        """lazily apply value to [start, stop)"""
         start = start_copy = start + self._size
         stop = stop_copy = stop + self._size
+        buffer_idx = self._buffer_idx
+
+        # For multi-field updates, initialize every lazy buffer field here.
+        self._lazy[buffer_idx] = value
+
         while start < stop:
             if start & 1:
-                self._lazy[start] += value
-                self.data[start] += value
+                self._apply_to_lazy(buffer_idx, start)
+                self._apply_to_data(buffer_idx, start)
                 start += 1
             if stop & 1:
                 stop -= 1
-                self._lazy[stop] += value
-                self.data[stop] += value
+                self._apply_to_lazy(buffer_idx, stop)
+                self._apply_to_data(buffer_idx, stop)
             start >>= 1
             stop >>= 1
 
-        # Tell all nodes above of the updated area of the updates
+        # Tell all nodes above the updated area of the updates.
         self._build(start_copy)
         self._build(stop_copy - 1)
 
+    add = apply
+
     def query(self, start, stop, default=0):
-        """func of data[start, stop)"""
+        """return func(data[start:stop])"""
         start += self._size
         stop += self._size
 
-        # Apply all the lazily stored queries
+        # Apply all the lazily stored updates.
         self._update(start)
         self._update(stop - 1)
 
-        res = default
+        buffer_idx = self._buffer_idx
+
+        # For multi-field data, initialize every data buffer field here.
+        self.data[buffer_idx] = default
         while start < stop:
             if start & 1:
-                res = self._func(res, self.data[start])
+                self._merge_data(buffer_idx, start, buffer_idx)
                 start += 1
             if stop & 1:
                 stop -= 1
-                res = self._func(res, self.data[stop])
+                self._merge_data(buffer_idx, stop, buffer_idx)
             start >>= 1
             stop >>= 1
-        return res
+
+        # For multi-field data, return every data buffer field here.
+        return self.data[buffer_idx]
 
     def __repr__(self):
         return "LazySegmentTree({0})".format(self.data)

tests/data_structures/test_LazySegmentTree.py

@@ -0,0 +1,76 @@
+import random
+
+import pytest
+
+from pyrival.data_structures.LazySegmentTree import LazySegmentTree
+
+
+@pytest.mark.parametrize("default,func", [(float("inf"), min), (float("-inf"), max)])
+def test_LazySegmentTree(default, func):
+    rng = random.Random(0)
+
+    for _ in range(100):
+        n = rng.randint(1, 100)
+        values = [rng.randint(-1000, 1000) for _ in range(n)]
+        tree = LazySegmentTree(values, default, func)
+
+        for _ in range(100):
+            start = rng.randrange(n)
+            stop = rng.randrange(start + 1, n + 1)
+            if rng.randrange(2):
+                value = rng.randint(-1000, 1000)
+                tree.add(start, stop, value)
+                for i in range(start, stop):
+                    values[i] += value
+            else:
+                assert tree.query(start, stop, default) == func(values[start:stop])
+
+
+class MinCountTree(LazySegmentTree):
+    """Example specialization: range chmin with counts, point query."""
+
+    INF = 10**9
+
+    def __init__(self, n):
+        super().__init__([(self.INF, 0)] * n, (self.INF, 0), self._merge)
+        self._lazy = [(self.INF, 0)] * len(self._lazy)
+
+    @staticmethod
+    def _merge(a, b):
+        if a[0] < b[0]:
+            return a
+        if b[0] < a[0]:
+            return b
+        return a[0], a[1] + b[1]
+
+    def _unset_lazy(self, idx):
+        self._lazy[idx] = self.INF, 0
+
+    def _apply_to_data(self, update_idx, data_idx):
+        self.data[data_idx] = self._merge(self._lazy[update_idx], self.data[data_idx])
+
+    def _apply_to_lazy(self, update_idx, lazy_idx):
+        self._lazy[lazy_idx] = self._merge(self._lazy[update_idx], self._lazy[lazy_idx])
+
+    def apply(self, start, stop, minimum, count):
+        super().apply(start, stop, (minimum, count))
+
+
+def test_LazySegmentTree_custom_min_count():
+    rng = random.Random(1)
+    n = 100
+    expected = [(MinCountTree.INF, 0) for _ in range(n)]
+    tree = MinCountTree(n)
+
+    for _ in range(1000):
+        start = rng.randrange(n)
+        stop = rng.randrange(start + 1, n + 1)
+        minimum = rng.randrange(20)
+        count = rng.randrange(1, 20)
+        tree.apply(start, stop, minimum, count)
+
+        for i in range(start, stop):
+            expected[i] = tree._merge(expected[i], (minimum, count))
+
+        idx = rng.randrange(n)
+        assert tree.query(idx, idx + 1, (MinCountTree.INF, 0)) == expected[idx]