diff --git a/pyrival/data_structures/LazySegmentTree.py b/pyrival/data_structures/LazySegmentTree.py index f6a5221..090ba96 100644 --- a/pyrival/data_structures/LazySegmentTree.py +++ b/pyrival/data_structures/LazySegmentTree.py @@ -1,4 +1,6 @@ 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 @@ -6,78 +8,111 @@ def __init__(self, data, default=0, func=max): 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) diff --git a/tests/data_structures/test_LazySegmentTree.py b/tests/data_structures/test_LazySegmentTree.py new file mode 100644 index 0000000..05150f7 --- /dev/null +++ b/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]