added Trie class, with tests

Author: jasonzhuyx

ml/misc/trie.py

@@ -0,0 +1,106 @@
+"""
+ml/misc/trie.py
+"""
+
+
+class Trie:
+    def __init__(self, words: list, name='trie'):
+        self._root = {}
+        _, self.count = Trie._build_tree(words, self._root)
+        self.name = name
+        pass
+
+    def __contains__(self, word):
+        found, node = self.has_match(word)
+        return found and None in node
+
+    def __iter__(self):
+        for result in Trie._iterator(self._root):
+            yield result
+
+    def __str__(self):
+        words = []
+        for word in self:
+            words.append(word)
+        return '\n'.join(words)
+
+    @staticmethod
+    def _add_node(tree, word, index=0):
+        if index >= len(word):
+            # adding a special key (None) to indicate the end of a word
+            tree[None] = '.'  # or tree[None] = word
+            return
+        char = word[index]
+        node = tree.get(char, {})
+        if node == {}:
+            tree[char] = node
+        Trie._add_node(node, word, index+1)
+        pass
+
+    @staticmethod
+    def _build_tree(words: list, tree: dict={}):
+        count = 0
+        for v in words:
+            value = v.strip()  # trimmed string value
+            if len(value) > 0:
+                # print('adding', value, 'to', tree)
+                Trie._add_node(tree, value)
+                count += 1
+        return tree, count
+
+    @staticmethod
+    def _iterator(tree, chars=[], prefix=''):
+        """
+        List all words from specific tree node.
+        """
+        for key in tree:
+            if key is not None:
+                chars.append(key)
+                for x in Trie._iterator(tree[key], chars, prefix):
+                    yield x
+                chars.pop()  # remove what has been processed
+            else:
+                suffix = ''.join(chars) if len(chars) > 0 else ''
+                complete_word = prefix + suffix
+                if complete_word:
+                    yield complete_word
+
+    @staticmethod
+    def _list(tree, chars=[], prefix=''):
+        return list(Trie._iterator(tree, chars, prefix))
+
+    def extend(self, words: list):
+        _, count = Trie._build_tree(words, self._root)
+        self.count += count
+        pass
+
+    def get_matches(self, prefix: str):
+        results = []
+        found, node = self.has_match(prefix)
+        if found and node is not None:
+            results.extend(Trie._iterator(node, [], prefix))
+        return results
+
+    def has(self, word):
+        return self.__contains__(word)
+
+    def has_match(self, prefix: str) -> (bool, dict):
+        """
+        Check if there is prefix in the trie list.
+        """
+        if not prefix:
+            return False, None
+        node = self._root
+        for ch in prefix:
+            tree = node.get(ch)
+            if tree is None:  # cannot find sub-tree/node for the char
+                return False, None
+            node = tree
+        return True, node
+
+    def has_prefix(self, prefix):
+        found, node = self.has_match(prefix)
+        return found and node is not None
+
+    def list(self):
+        return Trie._list(self._root)

ml/utils/domain_trie.py

@@ -1,6 +1,9 @@
 """
 DomainTrie class, as a data structure example.
 """
+from ml.utils.logger import get_logger
+
+LOGGER = get_logger(__name__)
 
 
 class DomainTrie(object):
@@ -10,7 +13,7 @@ class DomainTrie(object):
     """
 
     def __init__(self, domains):
-        self._trie = DomainTrie._make_tree(domains)
+        self._trie = DomainTrie._build_tree(domains)
 
     def __contains__(self, domain):
         """
@@ -19,21 +22,19 @@ def __contains__(self, domain):
         @param domain: A dot separated domain.
         @return boolean: True if domain is in the trie; otherwise, False.
         """
+        result = False
         parts = [part for part in domain.split('.')[::-1] if part != '']
-        res = False
-        ref = self._trie
+        node = self._trie  # start from the trie root
         for part in parts:
-            if part in ref:
-                ref = ref[part]
-                if '' in ref:
-                    res = True
-                    break
-            else:
+            if part not in node:
                 break
-        return res
+            node = node[part]
+            if '.' in node:
+                return True
+        return result
 
     def __iter__(self):
-        for result in DomainTrie._walk(self._trie, []):
+        for result in DomainTrie._list(self._trie, []):
             yield result
 
     def __str__(self):
@@ -43,37 +44,42 @@ def __str__(self):
         return str(result)
 
     @staticmethod
-    def _add_branch(domain, tree, idx=0):
-        # BASE CASE: out of character to add to this branch
-        if idx == len(domain):
-            tree[''] = None
-        else:
-            if '' not in tree:
-                domain_piece = domain[idx]
-                if domain_piece in tree:
-                    branch = tree[domain_piece]
-                else:
-                    branch = {}
-                    tree[domain_piece] = branch
-                DomainTrie._add_branch(domain, branch, idx + 1)
+    def _add_node(domain_pieces, tree, idx=0):
+        """
+        add specific level of domain part to the tree.
+        """
+        if idx >= len(domain_pieces):
+            tree['.'] = None
+            return
+        # stop at "." because it is only building topper levels
+        if "." not in tree:
+            part = domain_pieces[idx]
+            node = tree.get(part, {})
+            if node == {}:
+                tree[part] = node  # adding a new node
+            DomainTrie._add_node(domain_pieces, node, idx + 1)
 
     @staticmethod
-    def _make_tree(domains):
+    def _build_tree(domains):
         tree = {}
         for domain in domains:
-            domain_pieces = [piece for piece in domain.split('.') if piece != '']
+            d = domain.strip() if isinstance(domain, str) else ''
+            domain_pieces = [piece for piece in d.split('.') if piece != '']
             if len(domain_pieces) > 1:
-                DomainTrie._add_branch(domain_pieces[::-1], tree)
+                DomainTrie._add_node(domain_pieces[::-1], tree)
+            else:
+                LOGGER.warn('Invalid domain: "%s"', domain)
         return tree
 
     @staticmethod
-    def _walk(tree, domain_pieces):
-        result = []
-        if '' in tree:
-            result.append('.'.join(domain_pieces[::-1]))
-        else:
-            for piece in tree.keys():
+    def _list(tree, domain_pieces=[]):
+        results = []
+        if '.' in tree and len(domain_pieces) > 0:
+            results.append('.'.join(domain_pieces[::-1]))
+            return results
+        for piece in tree.keys():
+            if not piece == '.':
                 domain_pieces.append(piece)
-                result.extend(DomainTrie._walk(tree[piece], domain_pieces))
+                results.extend(DomainTrie._list(tree[piece], domain_pieces))
                 domain_pieces.pop()
-        return result
+        return results