diff --git a/wizard-work/beam/aho_corasick.py b/wizard-work/beam/aho_corasick.py new file mode 100644 index 0000000..bb82851 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/wizard-work/beam/test_aho_corasick.py b/wizard-work/beam/test_aho_corasick.py new file mode 100644 index 0000000..cf3ff67 --- /dev/null +++ b/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() \ No newline at end of file