beam: Aho-Corasick multi-pattern matcher with failure links and tests

wizard-work/beam/aho_corasick.py

@@ -0,0 +1,74 @@
+"""Aho-Corasick multi-pattern string matcher."""
+from __future__ import annotations
+
+from collections import deque
+from typing import List, Tuple
+
+
+class AhoCorasick:
+    def __init__(self, patterns: List[str]) -> None:
+        self.patterns = list(patterns)
+        self.goto: List[dict] = [{}]
+        self.fail: List[int] = [0]
+        self.direct_out: List[List[int]] = [[]]
+        for pi, pat in enumerate(patterns):
+            node = 0
+            for ch in pat:
+                nxt = self.goto[node].get(ch)
+                if nxt is None:
+                    nxt = len(self.goto)
+                    self.goto[node][ch] = nxt
+                    self.goto.append({})
+                    self.fail.append(0)
+                    self.direct_out.append([])
+                node = nxt
+            self.direct_out[node].append(pi)
+        self.out: List[List[int]] = [list(self.direct_out[i]) for i in range(len(self.goto))]
+        self._build_fail()
+
+    def _build_fail(self) -> None:
+        q: deque = deque()
+        for ch, nxt in list(self.goto[0].items()):
+            self.fail[nxt] = 0
+            self.out[nxt] = list(self.direct_out[nxt]) + list(self.out[0])
+            q.append(nxt)
+        while q:
+            r = q.popleft()
+            for ch, u in list(self.goto[r].items()):
+                q.append(u)
+                f = self.fail[r]
+                while f != 0 and ch not in self.goto[f]:
+                    f = self.fail[f]
+                self.fail[u] = self.goto[f].get(ch, 0)
+                if self.fail[u] == u:
+                    self.fail[u] = 0
+                self.out[u] = list(self.direct_out[u])
+                if self.fail[u] != 0:
+                    self.out[u].extend(self.out[self.fail[u]])
+
+    def search(self, text: str) -> List[Tuple[int, int]]:
+        res: List[Tuple[int, int]] = []
+        node = 0
+        n = len(text)
+        for i in range(n):
+            ch = text[i]
+            while node != 0 and ch not in self.goto[node]:
+                node = self.fail[node]
+            node = self.goto[node].get(ch, 0)
+            for pi in self.out[node]:
+                plen = len(self.patterns[pi])
+                if plen == 0:
+                    start = i
+                else:
+                    start = i + 1 - plen
+                res.append((start, pi))
+        # empty patterns also match at the end position
+        for pi, p in enumerate(self.patterns):
+            if p == "":
+                res.append((n, pi))
+        res.sort(key=lambda x: (x[0], x[1]))
+        return res
+
+
+def build(patterns):
+    return AhoCorasick(patterns)

wizard-work/beam/test_aho_corasick.py

@@ -0,0 +1,44 @@
+import unittest
+from aho_corasick import AhoCorasick
+
+
+class TestAhoCorasick(unittest.TestCase):
+    def test_basic(self):
+        ac = AhoCorasick(["he", "she", "his", "hers"])
+        r = ac.search("ushers")
+        # she@1, he@2, hers@2
+        self.assertIn((1, 1), r)  # she
+        self.assertIn((2, 0), r)  # he
+        self.assertIn((2, 3), r)  # hers
+
+    def test_overlap(self):
+        ac = AhoCorasick(["aa", "aaa"])
+        r = ac.search("aaaa")
+        # aa at 0,1,2 ; aaa at 0,1
+        self.assertEqual(sorted(r), sorted([(0,0),(1,0),(2,0),(0,1),(1,1)]))
+
+    def test_no_match(self):
+        ac = AhoCorasick(["xyz"])
+        self.assertEqual(ac.search("abc"), [])
+
+    def test_empty_pattern(self):
+        ac = AhoCorasick(["", "a"])
+        r = ac.search("a")
+        # empty matches at 0,1; 'a' at 0
+        self.assertEqual(sorted(r), sorted([(0,0),(1,0),(0,1)]))
+
+    def test_duplicate_patterns(self):
+        ac = AhoCorasick(["ab", "ab"])
+        r = ac.search("xabx")
+        self.assertEqual(r, [(1,0),(1,1)])
+
+    def test_nested(self):
+        ac = AhoCorasick(["a", "ab", "abc"])
+        r = ac.search("abc")
+        self.assertIn((0,0), r)
+        self.assertIn((0,1), r)
+        self.assertIn((0,2), r)
+
+
+if __name__ == "__main__":
+    unittest.main()