dict_trie.py 1.83 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26
def _hamming(path, node, word, distance):
    """
    :arg str path: Path taken so far to reach the current node.
    :arg dict node: Current node.
    :arg str word: Query word.
    :arg int distance: Amount of errors we can still make.

    :returns str:
    """
    if distance < 0:
        return ''
    if not word:
        if '' in node:
            return path
        return ''

    car, cdr = word[0], word[1:]
    for char in node:
        result = _hamming(
            path + char, node[char], cdr, distance - int(char != car))
        if result:
            return result

    return ''


27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
class Trie(object):
    def __init__(self, words):
        """
        Initialise the class.

        :arg list words: List of words.
        """
        self.root = {}

        self._build(words)

    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, None otherwise.
        """
        node = self.root

Jeroen F.J. Laros's avatar
Jeroen F.J. Laros committed
57 58
        for char in word:
            if char not in node:
59
                return {}
Jeroen F.J. Laros's avatar
Jeroen F.J. Laros committed
60
            node = node[char]
61 62 63 64

        return node

    def __contains__(self, word):
65
        return '' in self._find(word)
66 67 68 69

    def add(self, word):
        """
        Add a word to the trie.
Jeroen F.J. Laros's avatar
Jeroen F.J. Laros committed
70 71

        :arg str word: A word.
72 73 74
        """
        node = self.root

Jeroen F.J. Laros's avatar
Jeroen F.J. Laros committed
75 76 77 78
        for char in word:
            if char not in node:
                node[char] = {}
            node = node[char]
79 80 81 82

        node[''] = {}

    def has_prefix(self, word):
83
        return self._find(word) != {}
84 85 86

    def hamming(self, word, distance):
        return _hamming('', self.root, word, distance)