trie

Author: ivanmmarkovic

trie/trie.py

@@ -0,0 +1,111 @@
+
+class TrieNode:
+    def __init__(self, key, parent = None, children: dict = {}):
+        self.key = key
+        self.parent = parent
+        self.children:dict = {}
+        self.endchar: bool = False
+
+class Trie:
+    def __init__(self):
+        self.root: TrieNode = TrieNode(None)
+
+    def insert(self, string: str):
+        current: TrieNode = self.root
+        for character in string:
+            if character not in current.children:
+                current.children[character] = TrieNode(character, current)
+            current = current.children[character]
+        current.endchar = True
+
+    def contains(self, string: str)->bool:
+        current: TrieNode = self.root
+        for character in string:
+            if character not in current.children:
+                current = None
+                break
+            current = current.children[character]
+        if current is None:
+            return False
+        return current.endchar
+
+    def delete(self, string: str):
+        current: TrieNode = self.root
+        for character in string:
+            if character not in current.children:
+                current = None
+                break
+            current = current.children[character]
+        if current is None:
+            return
+        current.endchar = False
+        parent: TrieNode = current.parent
+        while parent is not None and not current.endchar and len(current.children) == 0:
+            del(parent.children[current.key])
+            current = parent
+            parent = current.parent
+
+    def prefix(self, prefix: str)->list:
+        current: TrieNode = self.root
+        for character in prefix:
+            if character not in current.children:
+                current = None
+                break
+            current = current.children[character]
+        if current is None:
+            return
+        words: list = []
+        self.helper(current, words, prefix)
+        return words
+
+    def helper(self, node: TrieNode, words: list, currentWord: str):
+        if node is None:
+            return
+        if node.endchar:
+            words.append(currentWord)
+        for key in node.children:
+            self.helper(node.children[key], words, currentWord + key)
+
+    def allWords(self)->list:
+        words: list = []
+        self.helper(self.root, words, "")
+        return words
+
+    def count(self)->int:
+        return self.countHelper(self.root)
+
+    def countHelper(self, node: TrieNode)->int:
+        if node is None:
+            return 0
+        sum: int = 0
+        if node.endchar:
+            sum += 1
+        for character in node.children:
+            sum += self.countHelper(node.children[character])
+        return sum
+
+trie = Trie()
+trie.insert("javascript")
+trie.insert("java")
+trie.insert("scala")
+trie.insert("scale")
+trie.insert("scalable")
+trie.insert("perl")
+
+print("Contains 'javascript' : ", trie.contains("javascript"))
+print("Contains 'java' : ", trie.contains("java"))
+print("Contains 'ruby' : ", trie.contains("ruby"))
+
+#trie.delete("java")
+trie.delete("javascript")
+
+
+print("Contains 'javascript' : ", trie.contains("javascript"))
+print("Contains 'java' : ", trie.contains("java"))
+print("Contains 'ruby' : ", trie.contains("ruby"))
+
+print(trie.prefix("scal")) # ['scala', 'scalable', 'scale']
+print(trie.prefix("java")) # ['java']
+
+print("All words", trie.allWords()) # All words ['java', 'scala', 'scalable', 'scale', 'perl']
+print("Count : ", trie.count())