Trie

pygorithm/data_structures/trie.py

@@ -0,0 +1,104 @@
+class Node:
+    def __init__(self, v, p=None, w=False):
+        self.word = w #If the node represents the end of a word or not
+        self.parent = p
+        self.value = v
+        self.children = {}
+
+
+class Trie:
+    def __init__(self):
+        self.root = Node('') #The root of the trie is always empty
+
+    def Insert(self, word):
+        """
+        Insert word in the trie. Starting from the root, move down the trie
+        following the path of characters in the word. If the nodes for the word
+        characters end, add them. When the last char is added, mark it as a
+        word-ending node.
+        """
+        l = len(word)
+        curr = self.root
+        for i, c in enumerate(word):
+            last = False
+            if(i == l-1):
+                #The last char of the word
+                last = True
+
+            if(c not in curr.children):
+                curr.children[c] = Node(c, curr, last)
+            elif(last):
+                #c already exists, but as it is the last char of word,
+                #it should now be flagged as a word in the trie.
+                curr.children[c].word = True
+
+            curr = curr.children[c]
+
+    def Search(self, word):
+        """
+        Searches for given word in trie. We want to find the last node for the
+        word. If we can't, then it means the word is not in the trie.
+        """
+        if self.FindFinalNode(word):
+            return True
+        else:
+            return False
+
+    def FindWords(self, prefix):
+        """
+        Find all words with the given prefix
+        """
+        v = self.FindFinalNode(prefix)
+        wList = self.BuildWordList(v, prefix)
+        if(v and v.word):
+            #v exists and the prefix is itself a word; add it to the list.
+            wList.append(prefix)
+
+        return wList
+
+    def FindFinalNode(self, word):
+        """
+        Returns the last node in given word. The process goes like this:
+        Start from the root. For every char in word, go down one level.
+        If we can't go down a level, then the word doesn't exist.
+        If we do, and the current char is the last char of the word and
+        the node we are currently at is a word, then we have found the given
+        word.
+        """
+        curr = self.root
+        l = len(word)
+
+        for i, c in enumerate(word):
+            if(c not in curr.children):
+                #There is no prefix of cWord + c
+                return None
+
+            if(i == l-1):
+                #Last char of word
+                return curr.children[c]
+
+            curr = curr.children[c]
+
+        return None
+
+    def BuildWordList(self, v, cWord):
+        """
+        Recursively builds the list of words.
+            * v: Node to check
+            * cWord : The word built up to v
+        """
+        if(not v):
+            return None
+
+        wList = []
+        for i, k in v.children.items():
+            tempWord = cWord + i
+
+            if(k.word):
+                #If the iterated prefix is a word, add it to the list
+                wList.append(tempWord)
+
+            #The list of words under tWord
+            wList.extend(self.BuildWordList(k, tempWord))
+
+        return wList

tests/test_data_structure.py

@@ -7,7 +7,8 @@
     linked_list,
     tree,
     graph,
-    heap)
+    heap,
+    trie)
 
 
 class TestStack(unittest.TestCase):
@@ -237,5 +238,22 @@ def test_heap(self):
         self.assertEqual(myHeap.queue, expectedResult)
 
 
+class TestTrie(unittest.TestCase):
+    def test_stack(self):
+        myTrie = trie.Trie()
+        myTrie.Insert('the')
+        myTrie.Insert('turtle')
+        myTrie.Insert('thesaurus')
+        myTrie.Insert('chocolate')
+        myTrie.Insert('flying')
+
+        self.assertEqual(myTrie.FindWords('th'), ['the', 'thesaurus'])
+        self.assertEqual(myTrie.FindWords('e'), None)
+
+        self.assertEqual(myTrie.Search('chocolate'), True)
+        self.assertEqual(myTrie.Search('flying'), True)
+        self.assertEqual(myTrie.Search('walking'), False)
+
+
 if __name__ == '__main__':
     unittest.main()