Create implement-trie-ii-prefix-tree.py

Author: kamyu104

Python/implement-trie-ii-prefix-tree.py

@@ -0,0 +1,73 @@
+# Time:  ctor:         O(1)
+#        insert:       O(n)
+#        count_word:   O(n)
+#        count_prefix: O(n)
+#        erase:        O(n)
+# Space: O(t), t is the number of nodes in trie
+
+class Node:
+    def __init__(self):
+        self.children = [None]*26
+        self.pcnt = 0
+        self.cnt = 0
+
+class Trie(object):
+
+    def __init__(self):
+        self.__trie = Node()
+
+    def insert(self, word):
+        """
+        :type word: str
+        :rtype: None
+        """
+        curr = self.__trie
+        curr.pcnt += 1
+        for c in word:
+            if curr.children[ord(c)-ord('a')] is None:
+                curr.children[ord(c)-ord('a')] = Node()
+            curr = curr.children[ord(c)-ord('a')]
+            curr.pcnt += 1
+        curr.cnt += 1
+
+    def countWordsEqualTo(self, word):
+        """
+        :type word: str
+        :rtype: int
+        """
+        curr = self.__trie
+        for c in word:
+            if curr.children[ord(c)-ord('a')] is None:
+                return 0
+            curr = curr.children[ord(c)-ord('a')]
+        return curr.cnt
+
+    def countWordsStartingWith(self, prefix):
+        """
+        :type prefix: str
+        :rtype: int
+        """
+        curr = self.__trie
+        for c in prefix:
+            if curr.children[ord(c)-ord('a')] is None:
+                return 0
+            curr = curr.children[ord(c)-ord('a')]
+        return curr.pcnt
+
+    def erase(self, word):
+        """
+        :type word: str
+        :rtype: None
+        """
+        cnt = self.countWordsEqualTo(word)
+        if not cnt:
+            return
+        curr = self.__trie
+        curr.pcnt -= 1
+        for c in word:
+            if curr.children[ord(c)-ord('a')].pcnt == 1:
+                curr.children[ord(c)-ord('a')] = None  # delete all unused nodes
+                return
+            curr = curr.children[ord(c)-ord('a')]
+            curr.pcnt -= 1
+        curr.cnt -= 1