Add compute_z_array Function with Unit Tests

algorithms/strings/__init__.py

@@ -38,4 +38,5 @@
 from .atbash_cipher import *
 from .longest_palindromic_substring import *
 from .knuth_morris_pratt import *
-from .panagram import *
+from .panagram import *
+from .z_algorithm import *

algorithms/strings/z_algorithm.py

@@ -0,0 +1,36 @@
+"""
+Algorithm that computes the Z-array for a given string using the Z Algorithm.
+
+The Z-array stores the length of the longest substring starting from index i which is also a prefix of the string.
+
+Time Complexity: O(n), where n is the length of the input string.
+"""
+
+def compute_z_array(input_string):
+    n = len(input_string)
+    z_array = [0] * n
+    left, right, k = 0, 0, 0
+
+    for i in range(1, n):
+        # Case 1: i is outside the current Z-box
+        if i > right:
+            left, right = i, i
+            # Expand the Z-box by matching with the prefix
+            while right < n and input_string[right] == input_string[right - left]:
+                right += 1
+            z_array[i] = right - left
+            right -= 1
+        else:
+            # Case 2: i is within the current Z-box
+            k = i - left
+            # If the Z[k] value is less than the remaining length of Z-box
+            if z_array[k] < right - i + 1:
+                z_array[i] = z_array[k]
+            else:
+                # Start from the right boundary and try to extend
+                left = i
+                while right < n and input_string[right] == input_string[right - left]:
+                    right += 1
+                z_array[i] = right - left
+                right -= 1
+    return z_array

tests/test_strings.py

@@ -42,7 +42,8 @@
     longest_palindrome,
     knuth_morris_pratt,
     panagram,
-    fizzbuzz
+    fizzbuzz,
+    compute_z_array
 )
 
 import unittest
@@ -728,5 +729,25 @@ def test_fizzbuzz(self):
         self.assertEqual(result, expected)
 
 
+class TestComputeZArray(unittest.TestCase):
+    """Test cases for compute_z_array function"""
+
+    def test_compute_z_array_basic(self):
+        self.assertEqual(compute_z_array("aaaaa"), [0, 4, 3, 2, 1])
+        self.assertEqual(compute_z_array("ababa"), [0, 0, 3, 0, 1])
+        self.assertEqual(compute_z_array("abcababc"), [0, 0, 0, 2, 0, 3, 0, 0])
+
+    def test_compute_z_array_no_prefix_match(self):
+        self.assertEqual(compute_z_array("abcde"), [0, 0, 0, 0, 0])
+
+    def test_compute_z_array_full_match(self):
+        self.assertEqual(compute_z_array("aaaa"), [0, 3, 2, 1])
+
+    def test_compute_z_array_single_character(self):
+        self.assertEqual(compute_z_array("a"), [0])
+
+    def test_compute_z_array_empty_string(self):
+        self.assertEqual(compute_z_array(""), [])
+
 if __name__ == "__main__":
     unittest.main()