beam: add Levenshtein edit distance with two-row DP and verified tests

wizard-work/beam/levenshtein.py

@@ -0,0 +1,86 @@
+"""Levenshtein edit distance with an optimal two-row DP.
+
+Provides:
+  edit_distance(a, b) -> int
+  ops(a, b) -> list of operations ('=', 'S', 'D', 'I') aligned to a
+  (substitution/deletion/insertion)
+
+Correctness: dp[i][j] = min edits to transform a[:i] into b[:j].
+Recurrence:
+  dp[0][j] = j
+  dp[i][0] = i
+  dp[i][j] = min(
+      dp[i-1][j] + 1,        # delete a[i-1]
+      dp[i][j-1] + 1,        # insert b[j-1]
+      dp[i-1][j-1] + (a[i-1] != b[j-1])  # match/substitute
+  )
+Space optimized to two rows: O(min(n,m)) memory, O(n*m) time.
+"""
+from __future__ import annotations
+
+
+def edit_distance(a: str, b: str) -> int:
+    if a == b:
+        return 0
+    if not a:
+        return len(b)
+    if not b:
+        return len(a)
+    # ensure b is the shorter to minimize memory
+    if len(a) < len(b):
+        a, b = b, a
+    # now a is longer-or-equal; iterate over b columns
+    prev = list(range(len(b) + 1))
+    for i, ca in enumerate(a, 1):
+        cur = [i] + [0] * len(b)
+        for j, cb in enumerate(b, 1):
+            cost = 0 if ca == cb else 1
+            cur[j] = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + cost)
+        prev = cur
+    return prev[-1]
+
+
+def ops(a: str, b: str) -> list[str]:
+    """Return operation tags aligned to characters of a: '=' match, 'S' sub,
+    'D' delete (a char absent in b), plus trailing 'I' inserts for extra b chars.
+    Reconstructs via backtracking on full table (O(n*m) memory) for clarity."""
+    n, m = len(a), len(b)
+    # full table
+    dp = [[0] * (m + 1) for _ in range(n + 1)]
+    for i in range(n + 1):
+        dp[i][0] = i
+    for j in range(m + 1):
+        dp[0][j] = j
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            cost = 0 if a[i - 1] == b[j - 1] else 1
+            dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j - 1] + 1, dp[i - 1][j - 1] + cost)
+    # backtrack
+    res: list[str] = []
+    i, j = n, m
+    while i > 0 or j > 0:
+        if i > 0 and j > 0:
+            cost = 0 if a[i - 1] == b[j - 1] else 1
+            if dp[i][j] == dp[i - 1][j - 1] + cost:
+                res.append("=" if cost == 0 else "S")
+                i -= 1
+                j -= 1
+                continue
+        if i > 0 and dp[i][j] == dp[i - 1][j] + 1:
+            res.append("D")
+            i -= 1
+            continue
+        # insert
+        res.append("I")
+        j -= 1
+    res.reverse()
+    # inserts that occur after all of a are already represented as trailing I's
+    return res
+
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) < 3:
+        print("usage: levenshtein.py STR1 STR2", file=sys.stderr)
+        sys.exit(2)
+    print(edit_distance(sys.argv[1], sys.argv[2]))

wizard-work/beam/test_levenshtein.py

@@ -0,0 +1,63 @@
+import random, string
+from levenshtein import edit_distance, ops
+
+
+def brute(a, b):
+    n, m = len(a), len(b)
+    dp = [[0]*(m+1) for _ in range(n+1)]
+    for i in range(n+1): dp[i][0]=i
+    for j in range(m+1): dp[0][j]=j
+    for i in range(1,n+1):
+        for j in range(1,m+1):
+            c = 0 if a[i-1]==b[j-1] else 1
+            dp[i][j]=min(dp[i-1][j]+1, dp[i][j-1]+1, dp[i-1][j-1]+c)
+    return dp[n][m]
+
+
+def test_known():
+    cases = [("kitten","sitting",3),("","abc",3),("abc","",3),
+             ("flaw","lawn",2),("same","same",0),("a","b",1),
+             ("","",-0 if False else 0)]
+    for a,b,e in cases:
+        assert edit_distance(a,b)==e, (a,b,edit_distance(a,b),e)
+
+
+def test_random_vs_brute():
+    rng = random.Random(42)
+    alpha = "ab"
+    for _ in range(2000):
+        la = rng.randint(0,8); lb = rng.randint(0,8)
+        a = "".join(rng.choice(alpha) for _ in range(la))
+        b = "".join(rng.choice(alpha) for _ in range(lb))
+        assert edit_distance(a,b)==brute(a,b), (a,b)
+
+
+def test_ops_roundtrip():
+    rng = random.Random(7)
+    for _ in range(500):
+        a = "".join(rng.choice("ab") for _ in range(rng.randint(0,6)))
+        b = "".join(rng.choice("ab") for _ in range(rng.randint(0,6)))
+        o = ops(a,b)
+        # apply ops to a -> should yield b
+        out = []
+        i = 0  # index into a
+        j = 0  # index into b
+        for tag in o:
+            if tag == "=":
+                assert a[i] == b[j]
+                out.append(a[i]); i += 1; j += 1
+            elif tag == "S":
+                out.append(b[j]); i += 1; j += 1
+            elif tag == "D":
+                i += 1
+            elif tag == "I":
+                # need current b char; find next unconsumed b
+                out.append(b[j]); j += 1
+        assert "".join(out) == b, (a,b,o,"".join(out))
+        edits = sum(1 for t in o if t in ("S","D","I"))
+        assert edits == edit_distance(a,b), (a,b,edits,edit_distance(a,b))
+
+
+if __name__ == "__main__":
+    test_known(); test_random_vs_brute(); test_ops_roundtrip()
+    print("ALL OK")