From c6ed0b223f5e55887b9c0e8b0e120f4b33ab0444 Mon Sep 17 00:00:00 2001 From: beam Date: Sun, 25 Oct 2026 16:08:22 +0000 Subject: [PATCH] beam: KMP substring search primitive with tests --- wizard-work/beam/kmp.py | 50 ++++++++++++++++++++++++++++++++++ wizard-work/beam/test_kmp.py | 52 ++++++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) create mode 100644 wizard-work/beam/kmp.py create mode 100644 wizard-work/beam/test_kmp.py diff --git a/wizard-work/beam/kmp.py b/wizard-work/beam/kmp.py new file mode 100644 index 0000000..4de90c3 --- /dev/null +++ b/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)) diff --git a/wizard-work/beam/test_kmp.py b/wizard-work/beam/test_kmp.py new file mode 100644 index 0000000..6450095 --- /dev/null +++ b/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()