beam: KMP substring search primitive with tests

wizard-work/beam/kmp.py

@@ -0,0 +1,50 @@
+"""Knuth-Morris-Pratt substring search.
+
+Provides:
+  - build_pi(pattern): prefix (failure) function, O(m)
+  - search(text, pattern): all start indices, O(n+m)
+  - search_count(text, pattern): number of (possibly overlapping) occurrences
+
+Correctness notes:
+  * Empty pattern matches at every position (0..len(text)) per Python str.find semantics? We return [] for empty pattern to avoid O(n) noise; callers wanting positions can handle separately. Actually we return list(range(len(text)+1)) to match str.find('') behaviour? Tests pin: empty pattern -> [] (documented).
+  * Overlapping occurrences are reported.
+"""
+
+from typing import List
+
+
+def build_pi(pattern: str) -> List[int]:
+    m = len(pattern)
+    pi = [0] * m
+    k = 0
+    for i in range(1, m):
+        while k > 0 and pattern[i] != pattern[k]:
+            k = pi[k - 1]
+        if pattern[i] == pattern[k]:
+            k += 1
+        pi[i] = k
+    return pi
+
+
+def search(text: str, pattern: str) -> List[int]:
+    if not pattern:
+        return []
+    n, m = len(text), len(pattern)
+    if m > n:
+        return []
+    pi = build_pi(pattern)
+    res: List[int] = []
+    k = 0
+    for i in range(n):
+        while k > 0 and text[i] != pattern[k]:
+            k = pi[k - 1]
+        if text[i] == pattern[k]:
+            k += 1
+        if k == m:
+            res.append(i - m + 1)
+            k = pi[k - 1]
+    return res
+
+
+def search_count(text: str, pattern: str) -> int:
+    return len(search(text, pattern))

wizard-work/beam/test_kmp.py

@@ -0,0 +1,52 @@
+import unittest
+from kmp import build_pi, search, search_count
+
+
+class TestKMP(unittest.TestCase):
+    def test_pi(self):
+        self.assertEqual(build_pi("aabaaab"), [0,1,0,1,2,2,3])
+        self.assertEqual(build_pi("aaaa"), [0,1,2,3])
+        self.assertEqual(build_pi(""), [])
+        self.assertEqual(build_pi("abc"), [0,0,0])
+
+    def test_search_basic(self):
+        self.assertEqual(search("ababcababa", "aba"), [0,5,7])
+        self.assertEqual(search("aaaa", "aa"), [0,1,2])
+        self.assertEqual(search("hello world", "xyz"), [])
+        self.assertEqual(search("abc", "abcd"), [])
+
+    def test_overlapping(self):
+        self.assertEqual(search("aaaaa", "aaa"), [0,1,2])
+
+    def test_single(self):
+        self.assertEqual(search("abc", "a"), [0])
+        self.assertEqual(search("abc", "c"), [2])
+
+    def test_empty_pattern(self):
+        self.assertEqual(search("abc", ""), [])
+        self.assertEqual(search("", ""), [])
+
+    def test_full_text(self):
+        self.assertEqual(search("abc", "abc"), [0])
+
+    def test_count(self):
+        self.assertEqual(search_count("ababab", "ab"), 3)
+
+    def test_unicode(self):
+        self.assertEqual(search("cafécafé", "café"), [0,4])
+
+    def test_random_vs_naive(self):
+        import random
+        rng = random.Random(42)
+        alpha = "ab"
+        for _ in range(200):
+            n = rng.randint(0, 60)
+            m = rng.randint(1, 8)
+            text = "".join(rng.choice(alpha) for _ in range(n))
+            pat = "".join(rng.choice(alpha) for _ in range(m))
+            naive = [i for i in range(len(text)-m+1) if text[i:i+m] == pat]
+            self.assertEqual(search(text, pat), naive)
+
+
+if __name__ == "__main__":
+    unittest.main()