Add lazy segment tree

pydatastructs/miscellaneous_data_structures/segment_tree.py

@@ -1,4 +1,5 @@
 from .stack import Stack
+from pydatastructs import OneDimensionalArray
 from pydatastructs.utils.misc_util import (TreeNode,
     Backend, raise_if_backend_is_not_python)
 
@@ -223,3 +224,138 @@ def query(self, start, end):
 
         return self._query(self._root, 0, len(self._array) - 1,
                            start, end)
+
+
+class OneDimensionalArraySegmentTreeLazy(OneDimensionalArraySegmentTree):
+
+    __slots__ = ["_func", "_array", "_lazy_node",
+                 "_neutral_element", "_root", "_backend"]
+
+    def __new__(cls, array, neutral_element, func, **kwargs):
+        backend = kwargs.get('backend', Backend.PYTHON)
+        raise_if_backend_is_not_python(cls, backend)
+
+        obj = object.__new__(cls)
+        obj._func = func
+        obj._array = array
+        obj._lazy_node = None
+        obj._neutral_element = neutral_element
+        obj._root = None
+        obj._backend = backend
+        return obj
+
+    @classmethod
+    def methods(self):
+        return ['__new__', 'build', 'update_range',
+                'update', 'query']
+
+    @property
+    def is_ready(self):
+        return self._root is not None
+
+    def build(self):
+        if self.is_ready:
+            return
+
+        recursion_stack = Stack(implementation='linked_list')
+        node = TreeNode((0, len(self._array) - 1), None, backend=self._backend)
+        lazy_node = TreeNode((0, len(self._array) - 1), None, backend=self._backend)
+        node.is_root = True
+        lazy_node.is_root = True
+        self._root = node
+        self._lazy_node = lazy_node
+        recursion_stack.push((node, lazy_node))
+
+        while not recursion_stack.is_empty:
+            node, lazy_node = recursion_stack.peek.key
+            start, end = node.key
+            if start == end:
+                node.data = self._array[start]
+                lazy_node.data = self._neutral_element
+                recursion_stack.pop()
+                continue
+
+            if (node.left is not None and
+                node.right is not None):
+                recursion_stack.pop()
+                node.data = self._func((node.left.data, node.right.data))
+                lazy_node.data = self._func((node.left.data, node.right.data))
+            else:
+                mid = (start + end) // 2
+                if node.left is None:
+                    left_node = TreeNode((start, mid), None)
+                    node.left = left_node
+                    lazy_left_node = TreeNode((start, mid), None)
+                    lazy_node.left = lazy_left_node
+                    recursion_stack.push((left_node, lazy_left_node))
+                if node.right is None:
+                    right_node = TreeNode((mid + 1, end), None)
+                    node.right = right_node
+                    lazy_rig_node = TreeNode((start, mid), None)
+                    lazy_node.right = lazy_rig_node
+                    recursion_stack.push((right_node, lazy_rig_node))
+
+    def _update_range(self, node, lazy_node, start, end, l, r, value):
+        if not self.is_ready:
+            raise ValueError("{} tree is not built yet. ".format(self) +
+                             "Call .build method to prepare the segment tree.")
+        if lazy_node.data != self._neutral_element:
+            node.data = self._func((node.data, lazy_node.data))
+            if start != end:
+                lazy_node.left.data = self._func((lazy_node.left.data, lazy_node.data))
+                lazy_node.right.data = self._func((lazy_node.right.data, lazy_node.data))
+            lazy_node.data = self._neutral_element
+        if r < start or end < l:
+            return
+
+        if l <= start and end <= r:
+            node.data = self._func((node.data, value))
+            if start != end:
+                lazy_node.left.data = self._func((lazy_node.left.data, value))
+                lazy_node.right.data = self._func((lazy_node.right.data, value))
+            return
+        mid = (start + end) // 2
+        self._update_range(node.left, lazy_node.left, start, mid, l, r, value)
+        self._update_range(node.right, lazy_node.right, mid + 1, end, l, r, value)
+        node.data = self._func((node.left.data, node.right.data))
+
+    def update_range(self, start, end, value):
+        if not self.is_ready:
+            raise ValueError("{} tree is not built yet. ".format(self) +
+                             "Call .build method to prepare the segment tree.")
+        self._update_range(self._root, self._lazy_node, 0, len(self._array) - 1,
+                           start, end, value)
+
+    def update(self, index, value):
+        if not self.is_ready:
+            raise ValueError("{} tree is not built yet. ".format(self) +
+                             "Call .build method to prepare the segment tree.")
+        self._update_range(self._root, self._lazy_node, 0, len(self._array) - 1,
+                           index, index, value)
+
+    def _query(self, node, lazy_node, start, end, l, r):
+        if r < start or end < l:
+            return None
+
+        if lazy_node.data != self._neutral_element:
+            node.data = self._func((node.data, lazy_node.data))
+            if start != end:
+                lazy_node.left.data = self._func((lazy_node.left.data, lazy_node.data))
+                lazy_node.right.data = self._func((lazy_node.right.data, lazy_node.data))
+            lazy_node.data = self._neutral_element
+
+        if l <= start and end <= r:
+            return node.data
+
+        mid = (start + end) // 2
+        left_result = self._query(node.left, lazy_node.left, start, mid, l, r)
+        right_result = self._query(node.right, lazy_node.right, mid + 1, end, l, r)
+        return self._func((left_result, right_result))
+
+    def query(self, start, end):
+        if not self.is_ready:
+            raise ValueError("{} tree is not built yet. ".format(self) +
+                             "Call .build method to prepare the segment tree.")
+
+        return self._query(self._root, 0, len(self._array) - 1,
+                           start, end)