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