Commit 5f6cdc3a authored by Jeroen F.J. Laros's avatar Jeroen F.J. Laros

Refactoring.

parent bbc2a94f
def _fill(node, alphabet, length): def _add(root, word):
"""Make a full trie using the characters in {alphabet}. """Add a word to the trie.
:arg dict node: Current node. :arg dict root: Root of the trie.
:arg tuple alphabet: Used alphabet. :arg str word: A word.
:arg int length: Length of the words to be generated. """
node = root
:returns iter: Trie containing all words of length {length} over alphabet for char in word:
{alphabet}. if char not in node:
node[char] = {}
node = node[char]
node[''] = {}
def _find(root, word):
"""Find the node after following the path in the trie given by {word}.
:arg dict root: Root of the trie.
:arg str word: A word.
:returns dict: The node if found, {} otherwise.
""" """
if not length: node = root
node[''] = {}
return for char in word:
if char not in node:
return {}
node = node[char]
for car in alphabet: return node
node[car] = {}
_fill(node[car], alphabet, length - 1)
def _remove(node, word): def _remove(node, word):
...@@ -43,7 +58,7 @@ def _remove(node, word): ...@@ -43,7 +58,7 @@ def _remove(node, word):
return result return result
def _to_list(path, node): def _iterate(path, node):
"""Convert a trie into a list. """Convert a trie into a list.
:arg str path: Path taken so far to reach the current node. :arg str path: Path taken so far to reach the current node.
...@@ -54,11 +69,30 @@ def _to_list(path, node): ...@@ -54,11 +69,30 @@ def _to_list(path, node):
if '' in node: if '' in node:
yield path yield path
for car in node: for char in node:
for result in _to_list(path + car, node[car]): for result in _iterate(path + char, node[char]):
yield result yield result
def _fill(node, alphabet, length):
"""Make a full trie using the characters in {alphabet}.
:arg dict node: Current node.
:arg tuple alphabet: Used alphabet.
:arg int length: Length of the words to be generated.
:returns iter: Trie containing all words of length {length} over alphabet
{alphabet}.
"""
if not length:
node[''] = {}
return
for char in alphabet:
node[char] = {}
_fill(node[char], alphabet, length - 1)
def _hamming(path, node, word, distance): def _hamming(path, node, word, distance):
"""Find all paths in the trie that are within a certain hamming distance of """Find all paths in the trie that are within a certain hamming distance of
{word}. {word}.
...@@ -129,61 +163,27 @@ class Trie(object): ...@@ -129,61 +163,27 @@ class Trie(object):
self.root = {} self.root = {}
if words: if words:
self._build(words) for word in words:
self.add(word)
def _build(self, words):
"""Build the trie.
:arg list words: List of words.
"""
for word in words:
self.add(word)
def _find(self, word):
"""Find the node after following the path in the trie given by {word}.
:arg str word: A word.
:returns dict: The node if found, {} otherwise.
"""
node = self.root
for char in word:
if char not in node:
return {}
node = node[char]
return node
def __contains__(self, word): def __contains__(self, word):
return '' in self._find(word) return '' in _find(self.root, word)
def add(self, word):
"""Add a word to the trie.
:arg str word: A word.
"""
node = self.root
for char in word: def __iter__(self):
if char not in node: return _iterate('', self.root)
node[char] = {}
node = node[char]
node[''] = {} def add(self, word):
_add(self.root, word)
def remove(self, word): def remove(self, word):
return _remove(self.root, word) return _remove(self.root, word)
def has_prefix(self, word): def has_prefix(self, word):
return self._find(word) != {} return _find(self.root, word) != {}
def fill(self, alphabet, length): def fill(self, alphabet, length):
_fill(self.root, alphabet, length) _fill(self.root, alphabet, length)
def to_list(self):
return _to_list('', self.root)
def all_hamming(self, word, distance): def all_hamming(self, word, distance):
return _hamming('', self.root, word, distance) return _hamming('', self.root, word, distance)
...@@ -201,7 +201,7 @@ class Trie(object): ...@@ -201,7 +201,7 @@ class Trie(object):
:returns str: Best match with {word}. :returns str: Best match with {word}.
""" """
if word in self: if _find(self.root, word):
return word return word
for i in range(1, distance + 1): for i in range(1, distance + 1):
...@@ -228,7 +228,7 @@ class Trie(object): ...@@ -228,7 +228,7 @@ class Trie(object):
:returns str: Best match with {word}. :returns str: Best match with {word}.
""" """
if word in self: if _find(self.root, word):
return word return word
for i in range(1, distance + 1): for i in range(1, distance + 1):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment