aho_corasick.py

from __future__ importannotations

fromcollectionsimportdeque


classAutomaton:
def__init__(self, keywords: list[str]):
self.adlist: list[dict] = []
self.adlist.append(
 {"value": "", "next_states": [], "fail_state": 0, "output": []}
 )

forkeywordinkeywords:
self.add_keyword(keyword)
self.set_fail_transitions()

deffind_next_state(self, current_state: int, char: str) ->int|None:
forstateinself.adlist[current_state]["next_states"]:
ifchar==self.adlist[state]["value"]:
returnstate
returnNone

defadd_keyword(self, keyword: str) ->None:
current_state=0
forcharacterinkeyword:
next_state=self.find_next_state(current_state, character)
ifnext_stateisNone:
self.adlist.append(
 {
"value": character,
"next_states": [],
"fail_state": 0,
"output": [],
 }
 )
self.adlist[current_state]["next_states"].append(len(self.adlist) -1)
current_state=len(self.adlist) -1
else:
current_state=next_state
self.adlist[current_state]["output"].append(keyword)

defset_fail_transitions(self) ->None:
q: deque=deque()
fornodeinself.adlist[0]["next_states"]:
q.append(node)
self.adlist[node]["fail_state"] =0
whileq:
r=q.popleft()
forchildinself.adlist[r]["next_states"]:
q.append(child)
state=self.adlist[r]["fail_state"]
while (
self.find_next_state(state, self.adlist[child]["value"]) isNone
andstate!=0
 ):
state=self.adlist[state]["fail_state"]
self.adlist[child]["fail_state"] =self.find_next_state(
state, self.adlist[child]["value"]
 )
ifself.adlist[child]["fail_state"] isNone:
self.adlist[child]["fail_state"] =0
self.adlist[child]["output"] = (
self.adlist[child]["output"]
+self.adlist[self.adlist[child]["fail_state"]]["output"]
 )

defsearch_in(self, string: str) ->dict[str, list[int]]:
"""
 >>> A = Automaton(["what", "hat", "ver", "er"])
 >>> A.search_in("whatever, err ... , wherever")
 {'what': [0], 'hat': [1], 'ver': [5, 25], 'er': [6, 10, 22, 26]}
 """
result: dict= {} # returns a dict with keywords and list of its occurrences
current_state=0
foriinrange(len(string)):
while (
self.find_next_state(current_state, string[i]) isNone
andcurrent_state!=0
 ):
current_state=self.adlist[current_state]["fail_state"]
next_state=self.find_next_state(current_state, string[i])
ifnext_stateisNone:
current_state=0
else:
current_state=next_state
forkeyinself.adlist[current_state]["output"]:
ifkeynotinresult:
result[key] = []
result[key].append(i-len(key) +1)
returnresult


if__name__=="__main__":
importdoctest

doctest.testmod()