Create segment_tree.py

Author: XenoBytesX

searches/segment_tree.py

@@ -0,0 +1,152 @@
+"""
+Author  : Sanjay Muthu <https://github.com/XenoBytesX>
+
+This is a Pure Python implementation of the Segment Tree Data Structure
+
+The problem statement is:
+    Given an array and q queries,
+    each query is one of two types:-
+        1. update:- (index, value)
+        update the array at index i to be equal to the new value
+        2. query:- (l, r)
+        print the result for the query from l to r
+        Here, the query depends on the problem which the segment tree is implemented on,
+        a common example of the query is sum or xor
+        (https://www.loginradius.com/blog/engineering/how-does-bitwise-xor-work/)
+
+Example:
+    array (a) = [5, 2, 3, 1, 7, 2, 9]
+    queries (q) = 2
+    query = sum
+
+    query 1:- update 1 3
+        - a[1] becomes 2
+        - a = [5, 3, 3, 1, 7, 2, 9]
+
+    query 2:- query 1 5
+        - a[1] + a[2] + a[3] + a[4] + a[5] = 3+3+1+7+2 = 16
+        - answer is 16
+
+Time Complexity:- O(N + Q)
+-- O(N) pre-calculation time to calculate the prefix sum array
+-- and O(1) time per each query = O(1 * Q) = O(Q) time
+
+Space Complexity:- O(N Log N + Q Log N)
+-- O(N Log N) time for building the segment tree
+-- O(log n) time for each query
+-- Q queries are there so total time complexity is O(Q Log n)
+
+Algorithm:-
+We first build the segment tree. An example of what the tree would look like:-
+(query type is sum)
+array = [5, 2, 3, 6, 1, 2]
+modified_array = [5, 2, 3, 6, 1, 2, 0, 0] size is 8 which a power of 2
+so we can build the segment tree
+
+segment tree:-
+                              19
+                            /   \
+                        /           \
+                    /                   \
+             /                               \
+            16                               3
+        /       \\                         /     \
+      /           \\                    /          \
+    /               \\               /               \
+    7               9               3               0
+  /   \\           /   \\           /   \\           /   \
+ /     \\         /     \\         /     \\         /     \
+/       \\       /       \\       /       \\       /       \
+5       2       3       6       1       2       0       0
+
+
+This segment tree cannot be stored in code so we convert it into a list
+
+segment tree list = [19, 16, 3, 7, 9, 3, 0, 5, 2, 3, 6, 1, 2, 0, 0]
+There is a property of this list that we can use to make the code much simpler
+segment tree list[2*i] and segment tree list[2*i+1]
+are the children of segment tree list[i]
+
+
+For Updating:-
+We first update the base element (the last row elements)
+and then slowly staircase up to update the entire segment tree part
+from the updated element
+
+For querying:-
+We start from the root(the topmost element) and go down, each node has one of 3 cases:-
+    Case 1. The node is completely inside the required range
+        then return the node value
+    Case 2. The node is completely outside the required range
+        then return 0
+    Case 3. The node is partially inside the required range
+        Query both the children and add their results and return that
+"""
+
+class SegmentTree:
+    def __init__(self, arr, merge_func, default):
+        """
+        Initializes the segment tree
+        :param arr: Input array
+        :param merge_func: The function which is used to merge
+        two elements of the segment tree
+        :param default: The default value for the nodes
+        (Ex:- 0 if merge_func is sum, inf if merge_func is min, etc.)
+        """
+        self.arr = arr
+        self.n = len(arr)
+
+        # while self.n is not a power of two
+        while (self.n & (self.n-1)) != 0:
+            self.n += 1
+            self.arr.append(default)
+
+        self.merge_func = merge_func
+        self.default = default
+        self.segment_tree = [default] * (2 * self.n)
+
+        for i in range(self.n):
+            self.segment_tree[self.n + i] = arr[i]
+
+        for i in range(self.n - 1, 0, -1):
+            self.segment_tree[i] = self.merge_func(self.segment_tree[2 * i],
+                                                   self.segment_tree[2 * i + 1])
+
+    def update(self, index, value):
+        """
+        Updates the value at an index and propagates the change to all parents
+        """
+        self.segment_tree[self.n + index] = value
+
+        while index >= 1:
+            index //= 2  # Go to the parent of index
+            self.segment_tree[index] = self.merge_func(self.segment_tree[2 * index],
+                                                       self.segment_tree[2 * index + 1])
+
+    def query(self, left, right, node_index=1, node_left=0, node_right=None):
+        """
+        Finds the answer of self.merge_query(left, left+1, left+2, left+3, ..., right)
+        """
+        if not node_right:
+            # We cant add self.n as the default value in the function
+            # because self itself is a parameter so we do it this way
+            node_right = self.n
+
+        # If the node is completely outside the query region we return the default value
+        if node_left > right or node_right < left:
+            return self.default
+
+        # If the node is completely inside the query region we return the node's value
+        if node_left > left and node_right < right:
+            return self.segment_tree[node_index]
+
+        # Else:-
+        # Find the middle element
+        mid = int((node_left + node_right) / 2)
+
+        # The answer is sum (or min or anything in the merge_func)
+        # of the query values of both the children nodes
+        return self.merge_func(
+            self.query(left, right, node_index * 2, node_left, mid),
+            self.query(left, right, node_index * 2 + 1, mid + 1, node_right)
+        )