74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
"""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) |