Repository: PyThaiNLP/pythainlp (python)

Issue:
Reduce Trie __init__ complexity
Current initialize code of Trie has a deep Node creation loop.

Propose to move most of them to a separated method to reduce code complexity. Was partially done in 30b4ad1.

- `add()` method will be added to `Trie`, this can be used at the initialization stage and also after the instance has been created (so a trie can has more word added later)
- `__len__` method will also be added, so a trie can use `len()` function to find its length (based on the length of internal set of words)


Pull request:
Refactor Trie __init__, reduce complexity
This PR will close #381 

Patch:
diff --git a/pythainlp/util/trie.py b/pythainlp/util/trie.py
index 445b04d9..7765ed16 100644
--- a/pythainlp/util/trie.py
+++ b/pythainlp/util/trie.py
@@ -14,25 +14,38 @@ class Trie:
         def __init__(self):
             self.end = False
             self.children = {}
-        
-        def add(self, ch: str):
-            child = self.children.get(ch)
-            if not child:
-                child = Trie.Node()
-                self.children[ch] = child
-            return child
 
     def __init__(self, words: Iterable[str]):
-        self.words = words
+        self.words = set(words)
         self.root = Trie.Node()
 
         for word in words:
-            cur = self.root
-            for ch in word:
-                cur = cur.add(ch)
-            cur.end = True
+            self.add(word)
+
+    def add(self, word: str) -> None:
+        """
+        Add a word to the trie.
+
+        :param str text: a word
+        """
+        self.words.add(word)
+        cur = self.root
+        for ch in word:
+            child = cur.children.get(ch)
+            if not child:
+                child = Trie.Node()
+                cur.children[ch] = child
+            cur = child
+        cur.end = True
 
     def prefixes(self, text: str) -> List[str]:
+        """
+        List all possible words from first sequence of characters in a word.
+
+        :param str text: a word
+        :return: a list of possible words
+        :rtype: List[str]
+        """
         res = []
         cur = self.root
         for i, ch in enumerate(text):
@@ -50,6 +63,9 @@ class Trie:
     def __iter__(self) -> Iterable[str]:
         yield from self.words
 
+    def __len__(self) -> int:
+        return len(self.words)
+
 
 def dict_trie(dict_source: Union[str, Iterable[str], Trie]) -> Trie:
     """


Test patch:
diff --git a/tests/test_util.py b/tests/test_util.py
index 84f4513e..2f02cea6 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -249,6 +249,12 @@ class TestUtilPackage(unittest.TestCase):
         self.assertIsNotNone(Trie(("ทอด", "ทอง", "ทาง")))
         self.assertIsNotNone(Trie(Trie(["ทดสอบ", "ทดลอง"])))
 
+        trie = Trie(["ทด", "ทดสอบ", "ทดลอง"])
+        self.assertTrue("ทด" in trie)
+        trie.add("ทบ")
+        self.assertEqual(len(trie), 4)
+        self.assertEqual(len(trie.prefixes("ทดสอบ")), 2)
+
         self.assertIsNotNone(dict_trie(Trie(["ลอง", "ลาก"])))
         self.assertIsNotNone(dict_trie(("ลอง", "สร้าง", "Trie", "ลน")))
         self.assertIsNotNone(dict_trie(["ลอง", "สร้าง", "Trie", "ลน"]))