suffix_array_lcp.py

divide_and_conquer/suffix_array_lcp.py

@@ -2,7 +2,6 @@
 
 import doctest
 
-
 def build_suffix_array(input_string: str) -> list[int]:
     """
     Build the suffix array for the given string.
@@ -14,7 +13,7 @@ def build_suffix_array(input_string: str) -> list[int]:
         list[int]: The suffix array (a list of starting indices of
                     suffixes in sorted order).
 
-    Example:
+    Examples:
         >>> build_suffix_array("banana")
         [5, 3, 1, 0, 4, 2]
     """
@@ -23,7 +22,6 @@ def build_suffix_array(input_string: str) -> list[int]:
     suffix_array = [suffix[1] for suffix in suffixes]
     return suffix_array
 
-
 def build_lcp_array(input_string: str, suffix_array: list[int]) -> list[int]:
     """
     Build the LCP array for the given string and suffix array.
@@ -35,9 +33,9 @@ def build_lcp_array(input_string: str, suffix_array: list[int]) -> list[int]:
     Returns:
         list[int]: The LCP array.
 
-    Example:
-        >>> suffix_arr = build_suffix_array("banana")
-        >>> build_lcp_array("banana", suffix_arr)
+    Examples:
+        >>> suffix_array = build_suffix_array("banana")
+        >>> build_lcp_array("banana", suffix_array)
         [0, 1, 3, 0, 0, 2]
     """
     n = len(input_string)
@@ -53,30 +51,29 @@ def build_lcp_array(input_string: str, suffix_array: list[int]) -> list[int]:
     for i in range(n):
         if rank[i] > 0:
             j = suffix_array[rank[i] - 1]
-            while (
-                (i + h < n)
-                and (j + h < n)
-                and (input_string[i + h] == input_string[j + h])
-            ):
+            while (i + h < n) and (j + h < n) and (input_string[i + h] == input_string[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__":
-    test_string = "banana"
-    suffix_array = build_suffix_array(test_string)
-    lcp_array = build_lcp_array(test_string, suffix_array)
+    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]}: {test_string[suffix_array[i]:]}")
+        print(f"{suffix_array[i]}: {s[suffix_array[i]:]}")
 
     print("\nLCP Array:")
     for i in range(1, len(lcp_array)):
-        print(
-            f"LCP between {test_string[suffix_array[i - 1]:]} and {test_string[suffix_array[i]]}: {lcp_array[i]}"
-        )
+        lcp_info = (f"LCP between {s[suffix_array[i - 1]:]} and "
+                    f"{s[suffix_array[i]]}: {lcp_array[i]}")
+        print(lcp_info)
+
+# Run doctests
+if __name__ == "__main__":
+    doctest.testmod()