Add Segment Tree Data Structure

Author: RishabhBhatnagar

Data_Structure/src/segment_tree.py

@@ -0,0 +1,92 @@
+from typing import TypeVar, List, Callable
+from functools import reduce
+
+
+T = TypeVar('T')  # Basic Template allowing any variable as input.
+
+
+def build(arr: List[T], n: int, func: Callable[[T, T], T]) -> List[T]:
+    """
+    builds the trees from the passed list.
+    Time Complexity: O(n)
+    
+    Parameters
+    ----------
+    arr: iterable
+       list of elements to be added to the data structure
+    func: Callable
+        function that will be binary operator
+
+    Returns
+    -------
+    Built Segment Tree
+    """
+    tree = [0 for i in range(n)]
+    tree.extend(arr)
+    for i in range(n-1, 0, -1):
+        tree[i] = func(tree[i << 1], tree[i << 1 | 1])
+    return tree
+
+
+def point_update(tree: List[T], n: int, pos: int, val: int, func: Callable[[T, T], T]) -> None:
+    """
+    Updates the tree with given value.
+    Time Complexity: O(lgn)
+    
+    Parameters
+    ----------
+    tree: list
+        built segment tree
+    n: int
+     size of segment tree
+    pos: int
+       index which is to be updated.
+    val: int
+        the value with which arr[pos] is to be updated.
+    func: Callable
+        binary operator which will be applied to the input parameters.
+    """
+    pos += n
+    tree[pos] = val
+    while pos > 1:
+        tree[pos >> 1] = func(tree[pos], tree[pos ^ 1])
+        pos >>= 1
+
+
+def query(tree: List[T], l: int, r: int, n: int, func: Callable[[T, T], T], start_ans=0, right_inclusive=True):
+    """
+    Parameters
+    ----------
+    tree: segment Tree
+    l: left
+    r: right
+    n: max size of array
+    func: Callable to be operated on the input parameters.
+    start_ans: seed value with which run is to be initialised
+    """
+
+    l += n
+    r += n + right_inclusive
+    ans = start_ans
+    while l < r:
+        if l & 1:
+            ans = func(ans, tree[l])
+            l += 1
+        if r & 1:
+            r -= 1
+            ans = func(ans, tree[r])
+        l >>= 1
+        r >>= 1
+    return ans
+
+
+
+if __name__ == '__main__':
+    arr = [3, 2, 4, 5, 6, 8, 2, 4, 5]
+    n = len(arr)
+    func = int.__mul__
+    tree = build(arr, n, func)
+    point_update(tree, n, 0, 0, func)
+    l, r = 2, 4
+    assert query(tree, l, r, n, func, start_ans=1) == reduce(func, arr[l: r+1]), "Segment Tree faulty"
+