From 335bf5b23acf19fe5c2b1b011579ab5d45782dd6 Mon Sep 17 00:00:00 2001 From: Mukundan314 Date: Wed, 6 Nov 2024 20:33:29 +0530 Subject: [PATCH 1/3] Refactor LazySegmentTree to be easier to modify --- pyrival/data_structures/LazySegmentTree.py | 76 ++++++++++++++-------- 1 file changed, 48 insertions(+), 28 deletions(-) diff --git a/pyrival/data_structures/LazySegmentTree.py b/pyrival/data_structures/LazySegmentTree.py index f6a5221b..842578fe 100644 --- a/pyrival/data_structures/LazySegmentTree.py +++ b/pyrival/data_structures/LazySegmentTree.py @@ -1,56 +1,74 @@ class LazySegmentTree: - def __init__(self, data, default=0, func=max): + def __init__(self, data): """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) + self.lazy = [0] * (2 * _size + 1) + self.data = [0] * (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]) + 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, a): + shift = self._size.bit_length() - a.bit_length() + return a << shift, (a << shift) + (1 << shift) + + def _unset_lazy(self, a): + """a: lazy_idx; unset a""" + self.lazy[a] = 0 + + def _apply_to_data(self, a, b): + """a: lazy_idx, b: data_idx; apply a to b""" + l, r = self._get_range(b) + self.data[b] += self.lazy[a] * (r - l) + + def _apply_to_lazy(self, a, b): + """a: lazy_idx, b: lazy_idx; apply a to b""" + self.lazy[b] += self.lazy[a] - self._lazy[2 * idx] += q - self._lazy[2 * idx + 1] += q - self.data[2 * idx] += q - self.data[2 * idx + 1] += q + def _merge_data(self, a, b, c): + """a: data_idx, b: data_idx, c: data_idx; merge a and b store result in c""" + self.data[c] = self.data[a] + self.data[b] 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) + _idx = idx >> i + self._apply_to_data(_idx, 2 * _idx) + self._apply_to_lazy(_idx, 2 * _idx) + self._apply_to_data(_idx, 2 * _idx + 1) + self._apply_to_lazy(_idx, 2 * _idx + 1) + self._unset_lazy(_idx) def _build(self, idx): """make the changes to idx be 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(2 * idx, 2 * 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 + + self.lazy[self._buffer_idx] = value + while start < stop: if start & 1: - self._lazy[start] += value - self.data[start] += value + self._apply_to_lazy(self._buffer_idx, start) + self._apply_to_data(self._buffer_idx, start) start += 1 if stop & 1: stop -= 1 - self._lazy[stop] += value - self.data[stop] += value + self._apply_to_lazy(self._buffer_idx, stop) + self._apply_to_data(self._buffer_idx, stop) start >>= 1 stop >>= 1 @@ -58,7 +76,7 @@ def add(self, start, stop, value): self._build(start_copy) self._build(stop_copy - 1) - def query(self, start, stop, default=0): + def query(self, start, stop): """func of data[start, stop)""" start += self._size stop += self._size @@ -67,17 +85,19 @@ def query(self, start, stop, default=0): self._update(start) self._update(stop - 1) - res = default + self.data[self._buffer_idx] = 0 + while start < stop: if start & 1: - res = self._func(res, self.data[start]) + self._merge_data(self._buffer_idx, start, self._buffer_idx) start += 1 if stop & 1: stop -= 1 - res = self._func(res, self.data[stop]) + self._merge_data(self._buffer_idx, stop, self._buffer_idx) start >>= 1 stop >>= 1 - return res + + return self.data[self._buffer_idx] def __repr__(self): return "LazySegmentTree({0})".format(self.data) From c5524ded01ce2a2828d002a4ae67f151e40da8c6 Mon Sep 17 00:00:00 2001 From: Cheran Date: Wed, 2 Sep 2026 10:26:19 +0530 Subject: [PATCH 2/3] Finish customizable LazySegmentTree refactor --- pyrival/data_structures/LazySegmentTree.py | 100 ++++++++++-------- tests/data_structures/test_LazySegmentTree.py | 76 +++++++++++++ 2 files changed, 130 insertions(+), 46 deletions(-) create mode 100644 tests/data_structures/test_LazySegmentTree.py diff --git a/pyrival/data_structures/LazySegmentTree.py b/pyrival/data_structures/LazySegmentTree.py index 842578fe..128cf311 100644 --- a/pyrival/data_structures/LazySegmentTree.py +++ b/pyrival/data_structures/LazySegmentTree.py @@ -1,55 +1,62 @@ class LazySegmentTree: - def __init__(self, data): + """Lazy segment tree whose update and merge operations can be specialized.""" + + 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._buffer_idx = 2 * _size - self.lazy = [0] * (2 * _size + 1) - self.data = [0] * (2 * _size + 1) + # Replace these arrays and the four methods below for custom 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)): + for i in range(_size - 1, 0, -1): self._merge_data(i + i, i + i + 1, i) def __len__(self): return self._len - def _get_range(self, a): - shift = self._size.bit_length() - a.bit_length() - return a << shift, (a << shift) + (1 << shift) + 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, a): - """a: lazy_idx; unset a""" - self.lazy[a] = 0 + def _unset_lazy(self, idx): + """clear the lazy update at idx""" + self._lazy[idx] = 0 - def _apply_to_data(self, a, b): - """a: lazy_idx, b: data_idx; apply a to b""" - l, r = self._get_range(b) - self.data[b] += self.lazy[a] * (r - l) + 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] - def _apply_to_lazy(self, a, b): - """a: lazy_idx, b: lazy_idx; apply a to b""" - self.lazy[b] += self.lazy[a] + 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, a, b, c): - """a: data_idx, b: data_idx, c: data_idx; merge a and b store result in c""" - self.data[c] = self.data[a] + self.data[b] + 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())): - _idx = idx >> i - self._apply_to_data(_idx, 2 * _idx) - self._apply_to_lazy(_idx, 2 * _idx) - self._apply_to_data(_idx, 2 * _idx + 1) - self._apply_to_lazy(_idx, 2 * _idx + 1) - self._unset_lazy(_idx) + """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._merge_data(2 * idx, 2 * idx + 1, idx) + self._merge_data(idx + idx, idx + idx + 1, idx) self._apply_to_data(idx, idx) idx >>= 1 @@ -57,47 +64,48 @@ def apply(self, start, stop, value): """lazily apply value to [start, stop)""" start = start_copy = start + self._size stop = stop_copy = stop + self._size - - self.lazy[self._buffer_idx] = value + buffer_idx = self._buffer_idx + self._lazy[buffer_idx] = value while start < stop: if start & 1: - self._apply_to_lazy(self._buffer_idx, start) - self._apply_to_data(self._buffer_idx, start) + self._apply_to_lazy(buffer_idx, start) + self._apply_to_data(buffer_idx, start) start += 1 if stop & 1: stop -= 1 - self._apply_to_lazy(self._buffer_idx, stop) - self._apply_to_data(self._buffer_idx, stop) + 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) - def query(self, start, stop): - """func of data[start, stop)""" + add = apply + + def query(self, start, stop, default=0): + """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) - self.data[self._buffer_idx] = 0 - + buffer_idx = self._buffer_idx + self.data[buffer_idx] = default while start < stop: if start & 1: - self._merge_data(self._buffer_idx, start, self._buffer_idx) + self._merge_data(buffer_idx, start, buffer_idx) start += 1 if stop & 1: stop -= 1 - self._merge_data(self._buffer_idx, stop, self._buffer_idx) + self._merge_data(buffer_idx, stop, buffer_idx) start >>= 1 stop >>= 1 - - return self.data[self._buffer_idx] + 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 00000000..05150f74 --- /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] From f950eac890fccbf560591c3e94cb5da8bbff1ad2 Mon Sep 17 00:00:00 2001 From: Cheran Date: Wed, 2 Sep 2026 10:37:47 +0530 Subject: [PATCH 3/3] Clarify LazySegmentTree customization contract --- pyrival/data_structures/LazySegmentTree.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/pyrival/data_structures/LazySegmentTree.py b/pyrival/data_structures/LazySegmentTree.py index 128cf311..090ba968 100644 --- a/pyrival/data_structures/LazySegmentTree.py +++ b/pyrival/data_structures/LazySegmentTree.py @@ -1,5 +1,5 @@ class LazySegmentTree: - """Lazy segment tree whose update and merge operations can be specialized.""" + """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""" @@ -10,7 +10,8 @@ def __init__(self, data, default=0, func=max): self._size = _size = 1 << (self._len - 1).bit_length() self._buffer_idx = 2 * _size - # Replace these arrays and the four methods below for custom operations. + # 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 @@ -65,6 +66,8 @@ def apply(self, start, stop, value): 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: @@ -95,6 +98,8 @@ def query(self, start, stop, default=0): self._update(stop - 1) 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: @@ -105,6 +110,8 @@ def query(self, start, stop, default=0): self._merge_data(buffer_idx, stop, buffer_idx) start >>= 1 stop >>= 1 + + # For multi-field data, return every data buffer field here. return self.data[buffer_idx] def __repr__(self):