Files
troupe-standup/wizard-work/beam/levenshtein.py
T

86 lines
2.6 KiB
Python

"""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]))