beam: add LCP array (Kasai) module with property tests

wizard-work/beam/lcp_array.py

@@ -0,0 +1,57 @@
+"""Longest common prefix (LCP) array for a suffix array.
+
+Given a text (str or bytes) and its suffix array (list of starting positions
+in ascending suffix order), compute the LCP array where lcp[i] is the length
+of the longest common prefix of suffix SA[i-1] and suffix SA[i], with
+lcp[0] = 0 by convention.
+
+Uses Kasai's O(n) algorithm: walk suffixes in text order, maintaining the
+current match length and decrementing by one each step.
+"""
+from __future__ import annotations
+
+
+def build_suffix_array(text):
+    """Naive O(n^2 log n) suffix array; fine for modest n."""
+    n = len(text)
+    return sorted(range(n), key=lambda i: text[i:])
+
+
+def build_lcp(text, sa):
+    n = len(text)
+    if len(sa) != n:
+        raise ValueError("suffix array length must equal text length")
+    if sorted(sa) != list(range(n)):
+        raise ValueError("invalid suffix array")
+    rank = [0] * n
+    for i, s in enumerate(sa):
+        rank[s] = i
+    lcp = [0] * n
+    k = 0
+    for i in range(n):
+        if rank[i] == 0:
+            k = 0
+            continue
+        j = sa[rank[i] - 1]
+        while i + k < n and j + k < n and text[i + k] == text[j + k]:
+            k += 1
+        lcp[rank[i]] = k
+        if k > 0:
+            k -= 1
+    return lcp
+
+
+def lcp_of(text, i, j):
+    """Direct LCP of suffixes starting at i and j."""
+    n = len(text)
+    k = 0
+    while i + k < n and j + k < n and text[i + k] == text[j + k]:
+        k += 1
+    return k
+
+
+if __name__ == "__main__":
+    t = "banana"
+    sa = build_suffix_array(t)
+    print(sa)
+    print(build_lcp(t, sa))

wizard-work/beam/test_lcp_array.py

@@ -0,0 +1,49 @@
+import os, sys, unittest
+sys.path.insert(0, os.path.dirname(__file__))
+from lcp_array import build_suffix_array, build_lcp, lcp_of
+
+
+class TestLCP(unittest.TestCase):
+    def _check(self, text):
+        sa = build_suffix_array(text)
+        lcp = build_lcp(text, sa)
+        if not text:
+            self.assertEqual(lcp, [])
+            return
+        self.assertEqual(lcp[0], 0)
+        for i in range(1, len(text)):
+            a, b = sa[i-1], sa[i]
+            self.assertEqual(lcp[i], lcp_of(text, a, b),
+                             "mismatch at %d for %r" % (i, text))
+
+    def test_banana(self):
+        self._check("banana")
+
+    def test_empty(self):
+        self.assertEqual(build_lcp("", []), [])
+
+    def test_single(self):
+        self._check("a")
+
+    def test_repeated(self):
+        self._check("aaaa")
+
+    def test_random(self):
+        import random
+        for _ in range(200):
+            n = random.randint(0, 12)
+            alphabet = "ab"
+            t = "".join(random.choice(alphabet) for _ in range(n))
+            self._check(t)
+
+    def test_invalid_sa_length(self):
+        with self.assertRaises(ValueError):
+            build_lcp("abc", [0, 1])
+
+    def test_invalid_sa_permutation(self):
+        with self.assertRaises(ValueError):
+            build_lcp("abc", [0, 0, 2])
+
+
+if __name__ == "__main__":
+    unittest.main()