|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +""" |
| 3 | + Aho-Corasick string search algorithm. |
| 4 | +
|
| 5 | + Author : Wojciech Muła, wojciech_mula@poczta.onet.pl |
| 6 | + WWW : http://0x80.pl |
| 7 | + License : public domain |
| 8 | +""" |
| 9 | + |
| 10 | +from collections import deque |
| 11 | + |
| 12 | +nil = object() # used to distinguish from None |
| 13 | + |
| 14 | +class TrieNode(object): |
| 15 | + """ |
| 16 | + Node of trie/Aho-Corasick automaton |
| 17 | + """ |
| 18 | + |
| 19 | + __slots__ = ['char', 'output', 'fail', 'children'] |
| 20 | + |
| 21 | + def __init__(self, char): |
| 22 | + """ |
| 23 | + Constructs an empty node |
| 24 | + """ |
| 25 | + |
| 26 | + self.char = char # character |
| 27 | + self.output = nil # an output function for this node |
| 28 | + self.fail = nil # fail link used by Aho-Corasick automaton |
| 29 | + self.children = {} # children |
| 30 | + |
| 31 | + def __repr__(self): |
| 32 | + """ |
| 33 | + Textual representation of node. |
| 34 | + """ |
| 35 | + |
| 36 | + if self.output is not nil: |
| 37 | + return "<TrieNode '%s' '%s'>" % (self.char, self.output) |
| 38 | + else: |
| 39 | + return "<TrieNode '%s'>" % self.char |
| 40 | + |
| 41 | + |
| 42 | +class Trie(object): |
| 43 | + """ |
| 44 | + Trie/Aho-Corasick automaton. |
| 45 | + """ |
| 46 | + |
| 47 | + def __init__(self): |
| 48 | + """ |
| 49 | + Construct an empty trie |
| 50 | + """ |
| 51 | + |
| 52 | + self.root = TrieNode('') |
| 53 | + |
| 54 | + |
| 55 | + def __get_node(self, word): |
| 56 | + """ |
| 57 | + Private function retrieving a final node of trie |
| 58 | + for given word |
| 59 | +
|
| 60 | + Returns node or None, if the trie doesn't contain the word. |
| 61 | + """ |
| 62 | + |
| 63 | + node = self.root |
| 64 | + for c in word: |
| 65 | + try: |
| 66 | + node = node.children[c] |
| 67 | + except KeyError: |
| 68 | + return None |
| 69 | + |
| 70 | + return node |
| 71 | + |
| 72 | + |
| 73 | + def get(self, word, default=nil): |
| 74 | + """ |
| 75 | + Retrieves output value associated with word. |
| 76 | +
|
| 77 | + If there is no word returns default value, |
| 78 | + and if default is not given rises KeyError. |
| 79 | + """ |
| 80 | + |
| 81 | + node = self.__get_node(word) |
| 82 | + output = nil |
| 83 | + if node: |
| 84 | + output = node.output |
| 85 | + |
| 86 | + if output is nil: |
| 87 | + if default is nil: |
| 88 | + raise KeyError("no key '%s'" % word) |
| 89 | + else: |
| 90 | + return default |
| 91 | + else: |
| 92 | + return output |
| 93 | + |
| 94 | + |
| 95 | + def keys(self): |
| 96 | + """ |
| 97 | + Generator returning all keys (i.e. word) stored in trie |
| 98 | + """ |
| 99 | + |
| 100 | + for key, _ in self.items(): |
| 101 | + yield key |
| 102 | + |
| 103 | + |
| 104 | + def values(self): |
| 105 | + """ |
| 106 | + Generator returning all values associated with words stored in a trie. |
| 107 | + """ |
| 108 | + |
| 109 | + for _, value in self.items(): |
| 110 | + yield value |
| 111 | + |
| 112 | + |
| 113 | + def items(self): |
| 114 | + """ |
| 115 | + Generator returning all keys and values stored in a trie. |
| 116 | + """ |
| 117 | + |
| 118 | + L = [] |
| 119 | + def aux(node, s): |
| 120 | + s = s + node.char |
| 121 | + if node.output is not nil: |
| 122 | + L.append((s, node.output)) |
| 123 | + |
| 124 | + for child in node.children.values(): |
| 125 | + if child is not node: |
| 126 | + aux(child, s) |
| 127 | + |
| 128 | + aux(self.root, '') |
| 129 | + return iter(L) |
| 130 | + |
| 131 | + |
| 132 | + def __len__(self): |
| 133 | + """ |
| 134 | + Calculates number of words in a trie. |
| 135 | + """ |
| 136 | + |
| 137 | + stack = deque() |
| 138 | + stack.append(self.root) |
| 139 | + n = 0 |
| 140 | + while stack: |
| 141 | + node = stack.pop() |
| 142 | + if node.output is not nil: |
| 143 | + n += 1 |
| 144 | + |
| 145 | + for child in node.children.values(): |
| 146 | + stack.append(child) |
| 147 | + |
| 148 | + return n |
| 149 | + |
| 150 | + |
| 151 | + def add_word(self, word, value): |
| 152 | + """ |
| 153 | + Adds word and associated value. |
| 154 | +
|
| 155 | + If word already exists, its value is replaced. |
| 156 | + """ |
| 157 | + if not word: |
| 158 | + return |
| 159 | + |
| 160 | + node = self.root |
| 161 | + for c in word: |
| 162 | + try: |
| 163 | + node = node.children[c] |
| 164 | + except KeyError: |
| 165 | + n = TrieNode(c) |
| 166 | + node.children[c] = n |
| 167 | + node = n |
| 168 | + |
| 169 | + node.output = value |
| 170 | + |
| 171 | + |
| 172 | + def clear(self): |
| 173 | + """ |
| 174 | + Clears trie. |
| 175 | + """ |
| 176 | + |
| 177 | + self.root = TrieNode('') |
| 178 | + |
| 179 | + |
| 180 | + def exists(self, word): |
| 181 | + """ |
| 182 | + Checks if whole word is present in the trie. |
| 183 | + """ |
| 184 | + |
| 185 | + node = self.__get_node(word) |
| 186 | + if node: |
| 187 | + return bool(node.output != nil) |
| 188 | + else: |
| 189 | + return False |
| 190 | + |
| 191 | + |
| 192 | + def match(self, word): |
| 193 | + """ |
| 194 | + Checks if word is a prefix of any existing word in the trie. |
| 195 | + """ |
| 196 | + |
| 197 | + return (self.__get_node(word) is not None) |
| 198 | + |
| 199 | + |
| 200 | + def make_automaton(self): |
| 201 | + """ |
| 202 | + Converts trie to Aho-Corasick automaton. |
| 203 | + """ |
| 204 | + |
| 205 | + queue = deque() |
| 206 | + |
| 207 | + # 1. |
| 208 | + for i in range(256): |
| 209 | + c = chr(i) |
| 210 | + if c in self.root.children: |
| 211 | + node = self.root.children[c] |
| 212 | + node.fail = self.root # f(s) = 0 |
| 213 | + queue.append(node) |
| 214 | + else: |
| 215 | + self.root.children[c] = self.root |
| 216 | + |
| 217 | + # 2. |
| 218 | + while queue: |
| 219 | + r = queue.popleft() |
| 220 | + for node in r.children.values(): |
| 221 | + queue.append(node) |
| 222 | + state = r.fail |
| 223 | + while node.char not in state.children: |
| 224 | + state = state.fail |
| 225 | + |
| 226 | + node.fail = state.children.get(node.char, self.root) |
| 227 | + |
| 228 | + |
| 229 | + def iter(self, string): |
| 230 | + """ |
| 231 | + Generator performs Aho-Corasick search string algorithm, yielding |
| 232 | + tuples containing two values: |
| 233 | + - position in string |
| 234 | + - outputs associated with matched strings |
| 235 | + """ |
| 236 | + state = self.root |
| 237 | + for index, c in enumerate(string): |
| 238 | + while c not in state.children: |
| 239 | + state = state.fail |
| 240 | + |
| 241 | + state = state.children.get(c, self.root) |
| 242 | + |
| 243 | + tmp = state |
| 244 | + output = [] |
| 245 | + while tmp is not nil: |
| 246 | + if tmp.output is not nil: |
| 247 | + output.append(tmp.output) |
| 248 | + |
| 249 | + tmp = tmp.fail |
| 250 | + |
| 251 | + if output: |
| 252 | + yield (index, output) |
| 253 | + |
| 254 | + def iter_long(self, string): |
| 255 | + """ |
| 256 | + Generator performs a modified Aho-Corasick search string algorithm, |
| 257 | + which maches only the longest word. |
| 258 | +
|
| 259 | + """ |
| 260 | + state = self.root |
| 261 | + last = None |
| 262 | + |
| 263 | + index = 0 |
| 264 | + while index < len(string): |
| 265 | + c = string[index] |
| 266 | + |
| 267 | + if c in state.children: |
| 268 | + state = state.children[c] |
| 269 | + |
| 270 | + if state.output is not nil: |
| 271 | + # save the last node on the path |
| 272 | + last = (state.output, index) |
| 273 | + |
| 274 | + index += 1 |
| 275 | + else: |
| 276 | + if last: |
| 277 | + # return the saved match |
| 278 | + yield last |
| 279 | + |
| 280 | + # and start over, as we don't want overlapped results |
| 281 | + # Note: this leads to quadratic complexity in the worst case |
| 282 | + index = last[1] + 1 |
| 283 | + state = self.root |
| 284 | + last = None |
| 285 | + else: |
| 286 | + # if no output, perform classic Aho-Corasick algorithm |
| 287 | + while c not in state.children: |
| 288 | + state = state.fail |
| 289 | + |
| 290 | + # corner case |
| 291 | + if last: |
| 292 | + yield last |
| 293 | + |
| 294 | + def find_all(self, string, callback): |
| 295 | + """ |
| 296 | + Wrapper on iter method, callback gets an iterator result |
| 297 | + """ |
| 298 | + for index, output in self.iter(string): |
| 299 | + callback(index, output) |
| 300 | + |
| 301 | + |
| 302 | + |
| 303 | +if __name__ == '__main__': |
| 304 | + |
| 305 | + def demo(): |
| 306 | + words = "he hers his she hi him man".split() |
| 307 | + |
| 308 | + t = Trie(); |
| 309 | + for w in words: |
| 310 | + t.add_word(w, w) |
| 311 | + |
| 312 | + s = "he rshershidamanza " |
| 313 | + |
| 314 | + t.make_automaton() |
| 315 | + for res in t.items(): |
| 316 | + print(res) |
| 317 | + |
| 318 | + for res in t.iter(s): |
| 319 | + print |
| 320 | + print('%s' % s) |
| 321 | + pos, matches = res |
| 322 | + for fragment in matches: |
| 323 | + print('%s%s' % ((pos - len(fragment) + 1) * ' ', fragment)) |
| 324 | + |
| 325 | + demo() |
| 326 | + |
| 327 | + |
| 328 | + def bug(): |
| 329 | + patterns = ['GT-C3303', 'SAMSUNG-GT-C3303K/'] |
| 330 | + text = 'SAMSUNG-GT-C3303i/1.0 NetFront/3.5 Profile/MIDP-2.0 Configuration/CLDC-1.1' |
| 331 | + |
| 332 | + t = Trie() |
| 333 | + for pattern in patterns: |
| 334 | + _ret = t.add_word(pattern, (0, pattern)) |
| 335 | + |
| 336 | + t.make_automaton() |
| 337 | + res = list(t.iter(text)) |
| 338 | + assert len(res) == 1, 'failed' |
| 339 | + bug() |
| 340 | + |
| 341 | +# vim: ts=4 sw=4 nowrap |
| 342 | + |
0 commit comments