Exemplo n.º 1
0
def load_corpus(corpus, case_sensitive=False):
    """Builds a Trie from a list of strings."""
    root = TrieNode()

    for word in corpus:
        if case_sensitive:
            root.insert(word, word)
        else:
            root.insert(word.lower(), word)

    return root
Exemplo n.º 2
0
class TestTrieNode(unittest.TestCase):
    def setUp(self):
        self.root = TrieNode()
        self.root.insert('ab', 'ab')
        self.root.insert('bc', 'bc')
        self.root.insert('cd', 'cd')
        self.root.insert('bcde', 'bcde')
        self.root.insert('bcdf', 'teste')

    def test_insert_root(self):
        self.assertEqual(len(self.root.children), 3)

    def test_children(self):
        node = self.root.children['a']
        self.assertEqual(len(node.children), 1)
        self.assertEqual(node.value, None)

        node = self.root.children['b']
        self.assertEqual(len(node.children), 1)
        self.assertEqual(node.value, None)

        node = node.children['c']
        self.assertEqual(len(node.children), 1)
        self.assertEqual(node.value, 'bc')

        node = node.children['d']
        self.assertEqual(len(node.children), 2)
        self.assertEqual(node.value, None)

        node_e = node.children['e']
        self.assertEqual(len(node_e.children), 0)
        self.assertEqual(node_e.value, 'bcde')

        node_f = node.children['f']
        self.assertEqual(len(node_f.children), 0)
        self.assertEqual(node_f.value, 'teste')

    def test_find_completions_node(self):
        node = find_completions_node(self.root, 'a')
        self.assertEqual(node, self.root.children['a'])

        node = find_completions_node(self.root, 'b')
        self.assertEqual(node, self.root.children['b'])

        node = find_completions_node(self.root, 'c')
        self.assertEqual(node, self.root.children['c'])

        node = find_completions_node(self.root, 'c')
        self.assertEqual(node, self.root.children['c'])

        node = find_completions_node(self.root, 'ab')
        self.assertEqual(node, self.root.children['a'].children['b'])

        node = find_completions_node(self.root, 'bc')
        self.assertEqual(node, self.root.children['b'].children['c'])

        node = find_completions_node(self.root, 'cd')
        self.assertEqual(node, self.root.children['c'].children['d'])

        node = find_completions_node(self.root, 'bcd')
        self.assertEqual(node,
                         self.root.children['b'].children['c'].children['d'])

        node = find_completions_node(self.root, 'bcde')
        self.assertEqual(
            node,
            self.root.children['b'].children['c'].children['d'].children['e'])

        node = find_completions_node(self.root, 'bcdf')
        self.assertEqual(
            node,
            self.root.children['b'].children['c'].children['d'].children['f'])