beam: add Z-algorithm utility + randomized brute-force tests

wizard-work/beam/test_z_algorithm.py

@@ -0,0 +1,46 @@
+import unittest
+from z_algorithm import z_function, search
+
+
+class TestZ(unittest.TestCase):
+    def test_empty(self):
+        self.assertEqual(z_function(""), [])
+
+    def test_single(self):
+        self.assertEqual(z_function("a"), [1])
+
+    def test_all_same(self):
+        self.assertEqual(z_function("aaaa"), [4, 3, 2, 1])
+
+    def test_ababa(self):
+        self.assertEqual(z_function("ababa"), [5, 0, 3, 0, 1])
+
+    def test_search_basic(self):
+        self.assertEqual(search("ab", "abxab"), [0, 3])
+
+    def test_search_overlap(self):
+        self.assertEqual(search("aa", "aaaa"), [0, 1, 2])
+
+    def test_search_empty_pattern(self):
+        self.assertEqual(search("", "abc"), [0, 1, 2, 3])
+
+    def test_search_none(self):
+        self.assertEqual(search("xyz", "abc"), [])
+
+    def test_search_full_match(self):
+        self.assertEqual(search("abc", "abc"), [0])
+
+    def test_random_brute(self):
+        import random, string
+        for _ in range(200):
+            n = random.randint(0, 40)
+            alpha = "ab"
+            s = "".join(random.choice(alpha) for _ in range(n))
+            p = "".join(random.choice(alpha) for _ in range(random.randint(0, 6)))
+            # brute
+            brute = [i for i in range(len(s) + 1) if s.startswith(p, i)] if p else list(range(len(s) + 1))
+            self.assertEqual(search(p, s), brute, (s, p))
+
+
+if __name__ == "__main__":
+    unittest.main()

wizard-work/beam/z_algorithm.py

@@ -0,0 +1,63 @@
+"""Z-algorithm: linear-time Z-array and substring search.
+
+Z[i] = length of the longest substring of s starting at i that is also a
+prefix of s.  Z[0] is conventionally len(s) (or 0, depending on source);
+here Z[0] = len(s) for convenience, matching the common competitive-programming
+convention.
+
+Complexity: O(n) time, O(n) space.
+"""
+
+
+def z_function(s: str):
+    """Return the Z-array of s."""
+    n = len(s)
+    z = [0] * n
+    if n == 0:
+        return z
+    z[0] = n
+    l = r = 0
+    for i in range(1, n):
+        if i < r:
+            z[i] = min(r - i, z[i - l])
+        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
+            z[i] += 1
+        if i + z[i] > r:
+            l, r = i, i + z[i]
+    return z
+
+
+def search(pattern: str, text: str):
+    """Return list of starting indices in text where pattern occurs (possibly overlapping)."""
+    if pattern == "":
+        return list(range(len(text) + 1))
+    if len(pattern) > len(text):
+        return []
+    combined = pattern + "\x00" + text
+    z = z_function(combined)
+    m = len(pattern)
+    res = []
+    offset = m + 1
+    for i in range(offset, len(combined)):
+        if z[i] >= m:
+            res.append(i - offset)
+    return res
+
+
+def _self_test():
+    assert z_function("") == []
+    assert z_function("a") == [1]
+    assert z_function("aaaa") == [4, 3, 2, 1]
+    assert z_function("ababa") == [5, 0, 3, 0, 1]
+    assert search("ab", "abxab") == [0, 3]
+    assert search("aa", "aaaa") == [0, 1, 2]
+    assert search("", "abc") == [0, 1, 2, 3]
+    assert search("xyz", "abc") == []
+    print("OK z_algorithm self-test")
+
+
+if __name__ == "__main__":
+    _self_test()
+    import sys
+    if len(sys.argv) > 1:
+        print(z_function(sys.argv[1]))