Aho-Corasick algorithm in Python
Builds a trie of the patterns with failure links and scans the text once, reporting all occurrences of all patterns in O(n + m + z) time.
- File
aho_corasick.py, 84 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/aho-corasick/. The authors are listed in prov/P3M1B3.jsonl.
from collections import deque
from collections.abc import Hashable, Sequence
ROOT = 0
NIL = -1
class Automaton:
"""The trie of the patterns with failure links and dictionary suffix links.
Node 0 is the root. goto[v] maps a character to a child, fail[v] is the
node of the longest proper suffix of str(v) that is also a node, dict_link[v]
is the nearest proper-suffix node that ends a pattern (NIL if none), and
out[v] lists the indices of the patterns equal to str(v).
"""
def __init__(self, patterns: Sequence[Sequence[Hashable]]) -> None:
self.lengths = [len(p) for p in patterns]
self.goto: list[dict[Hashable, int]] = [{}]
self.out: list[list[int]] = [[]]
for j, p in enumerate(patterns): # @line L1
v = ROOT
for a in p: # @line L2
if a not in self.goto[v]:
self.goto.append({})
self.out.append([])
self.goto[v][a] = len(self.goto) - 1
v = self.goto[v][a]
self.out[v].append(j) # @line L3
size = len(self.goto)
self.fail = [ROOT] * size
self.dict_link = [NIL] * size
queue: deque[int] = deque()
for u in self.goto[ROOT].values(): # @line L4
self.fail[u] = ROOT
queue.append(u)
while queue: # @line L5
v = queue.popleft() # @line L6
f = self.fail[v]
self.dict_link[v] = f if self.out[f] else self.dict_link[f] # @line L7
for a, u in self.goto[v].items(): # @line L8
f = self.fail[v] # @line L9
while f != ROOT and a not in self.goto[f]: # @line L10
f = self.fail[f] # @line L11
self.fail[u] = self.goto[f].get(a, ROOT) # @line L12
queue.append(u) # @line L13
def report(self, v: int, end: int, result: list[tuple[int, int]]) -> None:
"""Append (start, j) for each pattern j that is a suffix of str(v), longest first."""
w = v if self.out[v] else self.dict_link[v] # @line L19
while w != NIL: # @line L20
for j in self.out[w]: # @line L21
result.append((end - self.lengths[j], j))
w = self.dict_link[w] # @line L22
def search(self, text: Sequence[Hashable]) -> list[tuple[int, int]]:
"""All (start, j) with text[start:start + len(pattern j)] == pattern j."""
result: list[tuple[int, int]] = []
v = ROOT # @line L14
self.report(v, 0, result)
for i, c in enumerate(text): # @line L15
while v != ROOT and c not in self.goto[v]: # @line L16
v = self.fail[v]
v = self.goto[v].get(c, ROOT) # @line L17
self.report(v, i + 1, result) # @line L18
return result
def aho_corasick(
text: Sequence[Hashable], patterns: Sequence[Sequence[Hashable]]
) -> list[tuple[int, int]]:
"""Return every occurrence of every pattern in text as a pair (start, j).
Indices are 0-based, occurrences may overlap, and an empty pattern occurs at
each start 0, 1, ..., len(text). Pairs come in increasing order of end
position (start + length); at one end position, longer patterns come first,
and equal patterns in increasing index order. The time is O(n + m + z)
expected, for text length n, total pattern length m and z occurrences, with
O(m) extra space for the automaton.
"""
return Automaton(patterns).search(text)
test_aho_corasick.py 110 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/aho-corasick/. The authors are listed in prov/P3M1B3.jsonl.
import random
from aho_corasick import NIL, ROOT, Automaton, aho_corasick
def oracle(t: str, pats: list[str]) -> list[tuple[int, int]]:
return sorted(
(s, j) for j, p in enumerate(pats) for s in range(len(t) - len(p) + 1) if t[s : s + len(p)] == p
)
def node_strings(a: Automaton) -> list[str]:
s = [""] * len(a.goto)
stack = [ROOT]
while stack:
v = stack.pop()
for c, u in a.goto[v].items():
s[u] = s[v] + c
stack.append(u)
return s
def test_example() -> None:
pats = ["he", "she", "his", "hers"]
assert aho_corasick("ushers", pats) == [(1, 1), (2, 0), (2, 3)]
a = Automaton(pats)
s = node_strings(a)
fail = {s[v]: s[a.fail[v]] for v in range(1, len(s))}
assert fail == {"h": "", "s": "", "he": "", "hi": "", "sh": "h", "she": "he", "his": "s", "her": "", "hers": "s"}
dict_link = {s[v]: (s[a.dict_link[v]] if a.dict_link[v] != NIL else None) for v in range(1, len(s))}
assert dict_link["she"] == "he"
assert all(dict_link[x] is None for x in dict_link if x != "she")
def test_edges() -> None:
assert aho_corasick("", []) == []
assert aho_corasick("abc", []) == []
assert aho_corasick("", [""]) == [(0, 0)]
assert aho_corasick("ab", [""]) == [(0, 0), (1, 0), (2, 0)]
assert aho_corasick("ab", ["abc"]) == []
assert aho_corasick("aaaa", ["aa", "a"]) == [(0, 1), (0, 0), (1, 1), (1, 0), (2, 1), (2, 0), (3, 1)]
assert sorted(aho_corasick("abab", ["ab", "ab"])) == [(0, 0), (0, 1), (2, 0), (2, 1)]
def test_dictionary_links_needed() -> None:
# "b" is a suffix of "abc"'s prefix "ab" but str(v) = "ab" itself is not a pattern
assert sorted(aho_corasick("abd", ["abc", "b"])) == [(1, 1)]
# a pattern hidden two failure links deep
assert sorted(aho_corasick("xabcd", ["abcd", "bcd", "cd", "d", "bc"])) == oracle("xabcd", ["abcd", "bcd", "cd", "d", "bc"])
def test_fail_links_against_definition() -> None:
rng = random.Random(5)
for alphabet in ("a", "ab", "abcd"):
for _ in range(200):
pats = ["".join(rng.choice(alphabet) for _ in range(rng.randint(0, 7))) for _ in range(rng.randint(1, 5))]
a = Automaton(pats)
s = node_strings(a)
nodes = {x: v for v, x in enumerate(s)}
ends = {p for p in pats}
for v in range(1, len(s)):
x = s[v]
want = max((x[k:] for k in range(1, len(x) + 1) if x[k:] in nodes), key=len)
assert s[a.fail[v]] == want
dl = [x[k:] for k in range(1, len(x) + 1) if x[k:] in ends]
got = a.dict_link[v]
assert (s[got] if got != NIL else None) == (max(dl, key=len) if dl else None)
def random_patterns(rng: random.Random, alphabet: str, text: str) -> list[str]:
pats = []
for _ in range(rng.randint(0, 6)):
r = rng.random()
if r < 0.1:
pats.append("")
elif r < 0.4 and text:
i = rng.randrange(len(text))
pats.append(text[i : i + rng.randint(1, 6)])
elif r < 0.55 and pats:
p = rng.choice(pats)
pats.append(p[: rng.randint(0, len(p))] if rng.random() < 0.5 else p[rng.randint(0, len(p)) :])
elif r < 0.7:
unit = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 3)))
pats.append((unit * 5)[: rng.randint(1, 8)])
else:
pats.append("".join(rng.choice(alphabet) for _ in range(rng.randint(1, 6))))
return pats
def test_random_against_oracle() -> None:
rng = random.Random(2024)
for alphabet in ("a", "ab", "abcd"):
for _ in range(700):
if rng.random() < 0.3:
unit = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 3)))
t = (unit * 20)[: rng.randint(0, 30)]
else:
t = "".join(rng.choice(alphabet) for _ in range(rng.randint(0, 30)))
pats = random_patterns(rng, alphabet, t)
got = aho_corasick(t, pats)
assert sorted(got) == oracle(t, pats)
ends = [s + len(pats[j]) for s, j in got]
assert ends == sorted(ends)
def test_sequences_of_integers() -> None:
assert sorted(aho_corasick([1, 2, 1, 2, 1], [[1, 2, 1], [2]])) == [(0, 0), (1, 1), (2, 0), (3, 1)]
What the next level needs
- the tests did not run in this build