dict_trie.py 1.15 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
27
28
29
30
31
32
33
34
35
36
37
38
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

        for character in word:
            if character not in node:
                return {}
            node = node[character]

        return node

    def __contains__(self, word):
Laros's avatar
Laros committed
39
        return '' in self._find(word)
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

    def add(self, word):
        """
        Add a word to the trie.
        """
        node = self.root

        for character in word:
            if character not in node:
                node[character] = {}
            node = node[character]

        node[''] = {}

    def has_prefix(self, word):
Laros's avatar
Laros committed
55
        return self._find(word) != {}