Persistent Segment Tree

data_structures/persistent_segment_tree.py

@@ -0,0 +1,151 @@
+class Node:
+    def __init__(self, value=0):
+        """
+        Initialize a segment tree node.
+
+        Args:
+            value (int): The value of the node.
+        """
+        self.value = value
+        self.left = None
+        self.right = None
+
+
+class PersistentSegmentTree:
+    def __init__(self, arr):
+        """
+        Initialize the persistent segment tree with the given array.
+
+        Args:
+            arr (list): The initial array to build the segment tree.
+        """
+        self.n = len(arr)
+        self.roots = []
+        self.roots.append(self._build(arr, 0, self.n - 1))
+
+    def _build(self, arr, start, end):
+        """
+        Recursively build the segment tree.
+
+        Args:
+            arr (list): The input array.
+            start (int): The starting index of the segment.
+            end (int): The ending index of the segment.
+
+        Returns:
+            Node: The root node of the segment tree for the current segment.
+        """
+        if start == end:
+            return Node(arr[start])
+
+        mid = (start + end) // 2
+        node = Node()
+        node.left = self._build(arr, start, mid)
+        node.right = self._build(arr, mid + 1, end)
+        node.value = node.left.value + node.right.value
+        return node
+
+    def update(self, version, index, value):
+        """
+        Update the value at the specified index in the specified version.
+
+        Args:
+            version (int): The version of the segment tree to update.
+            index (int): The index to update.
+            value (int): The new value to set at the index.
+
+        Returns:
+            int: The index of the new version of the root node.
+        """
+        new_root = self._update(self.roots[version], 0, self.n - 1, index, value)
+        self.roots.append(new_root)
+        return len(self.roots) - 1  # return the index of the new version
+
+    def _update(self, node, start, end, index, value):
+        """
+        Recursively update the segment tree.
+
+        Args:
+            node (Node): The current node of the segment tree.
+            start (int): The starting index of the segment.
+            end (int): The ending index of the segment.
+            index (int): The index to update.
+            value (int): The new value to set at the index.
+
+        Returns:
+            Node: The new root node after the update.
+        """
+        if start == end:
+            new_node = Node(value)
+            return new_node
+
+        mid = (start + end) // 2
+        new_node = Node()
+        if index <= mid:
+            new_node.left = self._update(node.left, start, mid, index, value)
+            new_node.right = node.right
+        else:
+            new_node.left = node.left
+            new_node.right = self._update(node.right, mid + 1, end, index, value)
+
+        new_node.value = new_node.left.value + new_node.right.value
+        return new_node
+
+    def query(self, version, left, right):
+        """
+        Query the sum of values in the range [left, right] for the specified version.
+
+        Args:
+            version (int): The version of the segment tree to query.
+            left (int): The left index of the range.
+            right (int): The right index of the range.
+
+        Returns:
+            int: The sum of the values in the specified range.
+        """
+        return self._query(self.roots[version], 0, self.n - 1, left, right)
+
+    def _query(self, node, start, end, left, right):
+        """
+        Recursively query the segment tree.
+
+        Args:
+            node (Node): The current node of the segment tree.
+            start (int): The starting index of the segment.
+            end (int): The ending index of the segment.
+            left (int): The left index of the range.
+            right (int): The right index of the range.
+
+        Returns:
+            int: The sum of the values in the specified range.
+        """
+        if right < start or end < left:
+            return 0  # out of range
+
+        if left <= start and end <= right:
+            return node.value  # completely within range
+
+        mid = (start + end) // 2
+        sum_left = self._query(node.left, start, mid, left, right)
+        sum_right = self._query(node.right, mid + 1, end, left, right)
+        return sum_left + sum_right
+
+
+# Example usage and doctests
+if __name__ == "__main__":
+    import doctest
+
+    # Creating an initial array
+    arr = [1, 2, 3, 4, 5]
+    pst = PersistentSegmentTree(arr)
+
+    # Querying the initial version
+    assert pst.query(0, 0, 4) == 15  # sum of [1, 2, 3, 4, 5]
+
+    # Updating index 2 to value 10 in version 0
+    new_version = pst.update(0, 2, 10)
+
+    # Querying the updated version
+    assert pst.query(new_version, 0, 4) == 22  # sum of [1, 2, 10, 4, 5]
+    assert pst.query(0, 0, 4) == 15  # original version unchanged
+

divide_and_conquer/suffix_array_lcp.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+
+def build_suffix_array(s: str) -> list[int]:
+    """
+    Build the suffix array for the given string.
+
+    Parameters:
+        s (str): The input string.
+
+    Returns:
+        list[int]: The suffix array (a list of starting indices of
+                    suffixes in sorted order).
+    """
+    suffixes = [(s[i:], i) for i in range(len(s))]
+    suffixes.sort()  # Sort the suffixes lexicographically
+    suffix_array = [suffix[1] for suffix in suffixes]
+    return suffix_array
+
+def build_lcp_array(s: str, suffix_array: list[int]) -> list[int]:
+    """
+    Build the LCP array for the given string and suffix array.
+
+    Parameters:
+        s (str): The input string.
+        suffix_array (list[int]): The suffix array.
+
+    Returns:
+        list[int]: The LCP array.
+    """
+    n = len(s)
+    rank = [0] * n
+    lcp = [0] * n
+
+    # Compute the rank of each suffix
+    for i, suffix_index in enumerate(suffix_array):
+        rank[suffix_index] = i
+
+    # Compute the LCP array
+    h = 0
+    for i in range(n):
+        if rank[i] > 0:
+            j = suffix_array[rank[i] - 1]
+            while (i + h < n) and (j + h < n) and (s[i + h] == s[j + h]):
+                h += 1
+            lcp[rank[i]] = h
+            if h > 0:
+                h -= 1  # Decrease h for the next suffix
+    return lcp
+
+# Example usage
+if __name__ == "__main__":
+    s = "banana"
+    suffix_array = build_suffix_array(s)
+    lcp_array = build_lcp_array(s, suffix_array)
+
+    print("Suffix Array:")
+    for i in range(len(suffix_array)):
+        print(f"{suffix_array[i]}: {s[suffix_array[i]:]}")
+
+    print("\nLCP Array:")
+    for i in range(1, len(lcp_array)):
+        print(f"LCP between {s[suffix_array[i - 1]:]} and {s[suffix_array[i]]}: {lcp_array[i]}")