from collections import deque
class State:
sid = None ## store the id of state
value = None ## stores values of state
tranList = None ## used to store the list of next states for transition
outputSet = None ## it is set datastructure for storing the outputs at that state
failState = None
def __init__(self ,sid, val):
self.sid = sid
self.value = val
self.tranList = []
self.failState = 0
self.outputSet = set()
def getTransition(self, val):
""" this function gets the next state on input val"""
for node in self.tranList:
if node.value == val:
return node
return None
def testTransition(self, val):
""" This checks whether there is transition or not on input val"""
""" for current state, the transition is always true on any input"""
if self.sid == 0:
return True
else:
for nd in self.tranList:
if nd.value == val:
return True
return False
def addOutput(self, key):
"""This adds the key to the output in the state"""
self.outputSet = self.outputSet ^ key
##------------------------------------------------------------------------
class ahoCorasick:
root = None
newstate = None
def __init__(self):
self.root = State(0, ' ')
self.newstate = 0
def addKeyword(self, keywords):
"""Adds the keyword in the tree"""
for key in keywords.split(' '):
j = 0
state = 0
current = self.root
key = key.upper()
while j < len(key):
ch = key[j]
j = j+ 1
child = current.getTransition(ch)
if child != None:
current = child
else:
self.newstate = self.newstate +1
nd = State(self.newstate, ch)
current.tranList.append(nd)
current = nd
while j < len(key):
self.newstate = self.newstate +1
nd2 = State(self.newstate, key[j])
current.tranList.append(nd2)
current = nd2
j = j+1
break
current.outputSet.add(key)
##-------------------------------------------------------------------
def setFailTransitions(self):
"""Sets the fail transitions in tree"""
queue = deque()
current = self.root
child = self.root
for nd in self.root.tranList:
queue.append(nd)
nd.failState = self.root
while len(queue) != 0:
r = queue.popleft()
for nd in r.tranList:
queue.append(nd)
state = r.failState
val = nd.value
current = state
while True:
if current.testTransition(val) == False:
current = current.failState
else:
break
child = current.getTransition(val)
if child == None:
nd.failState = current
else:
nd.failState = child
nd.addOutput(nd.failState.outputSet)
##--------------------------------------------------------------------------------------------------
def findSubstrings(self, findStr):
""" Finds all substrings of input which are keywords in the tree"""
for string in findStr.split(' '):
string = string.upper()
print "Finding substrings in ", string
current = self.root
j = 0
while j < len(string):
while True:
if current.testTransition(string[j]) == False:
current = current.failState
else:
child = current.getTransition(string[j])
## print "before break", child.sid
break
if child != None:
## print "in none"
current = child
if len(child.outputSet) != 0:
print j
itr = iter(child.outputSet)
for keyw in itr:
print keyw
j = j + 1
##---------------------------------------------------------
def displayTree(self):
""" It is used to display the tree of keywords. Prints ID of node and value of node"""
queue = deque()
for nd in self.root.tranList:
queue.append(nd)
while len(queue) !=0:
node = queue.popleft()
for nd in node.tranList:
queue.append(nd)
print node.sid, node.value
def displayOutput(self):
""" This function displays the outputs at a state"""
queue = deque()
for nd in self.root.tranList:
queue.append(nd)
while len(queue) !=0:
node = queue.popleft()
for nd in node.tranList:
queue.append(nd)
itr = iter(node.outputSet)
if len(node.outputSet) !=0:
print node.sid
for string in itr:
print string
if (__name__ == "__main__"):
x = ahoCorasick()
""" Usage: Create object of ahoCorasick
to enter keywords use addKeyword("string of keywords")
then call setFailTransitions (fail function)
to find substrings of string use findSubstrings"""
## x.addKeyword("he")
## x.addKeyword("she")
## x.addKeyword("his")
## x.addKeyword("hers")
## x.addKeyword("ATC")
## x.addKeyword("TC")
## x.displayOutput()
#### x.enter("help")
#### x.enter("hi")
x.addKeyword("john jane")
x.setFailTransitions()
## x.findSubstrings("ACGATCTCTCGATC")
x.findSubstrings("johnjane")
