Merge pull request prakhar1989#18 from Xuefeng-Zhu/patch-11

prakhar1989 · prakhar1989 · commit c5e7f803b83a · 2015-05-15T09:11:52.000+03:00
Improve trie
diff --git a/trees/trie.py b/trees/trie.py
@@ -1,85 +1,100 @@
-""" Tries in python 
+""" Tries in python
 Methods -  insert_key(k, v)
            has_key(k)
            retrie_val(k)
            start_with_prefix(prefix)
 """
-# HELPERS #
-def _get_child_branches(tr):
-    if tr == []:
-        return []
-    return tr[1:]
 
-def _get_child_branch(tr, c):
-    for branch in _get_child_branches(tr):
+
+def _get_child_branches(trie):
+    """
+    Helper method for getting branches
+    """
+    return trie[1:]
+
+
+def _get_child_branch(trie, c):
+    """
+    Get branch matching the character
+    """
+    for branch in _get_child_branches(trie):
         if branch[0] == c:
             return branch
+
     return None
 
-def _retrive_branch(k, trie_list):
-    if k == "":
+
+def _retrive_branch(k, trie):
+    """
+    Get branch matching the key word
+    """
+    if not k:
         return None
-    tr = trie_list
+
     for c in k:
-        child_branch = _get_child_branch(tr, c)
+        child_branch = _get_child_branch(trie, c)
         if not child_branch:
             return None
-        tr = child_branch
-    return tr
+        trie = child_branch
+
+    return trie
+
 
 def _is_trie_bucket(bucket):
     if len(bucket) != 2:
         return False
-    if type(bucket[1]) is tuple:
-        return True
+
+    return type(bucket[1]) is tuple
+
 
 def _get_bucket_key(bucket):
     if not _is_trie_bucket(bucket):
         return None
-    return bucket[1][0] 
 
-# HAS_KEY #
-def has_key(k, tr):
-    if k == "":
-        return None
-    key_tuple = _retrive_branch(k, tr)
-    if not key_tuple:
-        return False
-    return True
+    return bucket[1][0]
 
-# RETRIE_VAL
-def retrie_val(k, tr):
-    if k == "":
-        return None
-    key_tuple = _retrive_branch(k, tr)
+
+def has_key(k, trie):
+    """
+    Check if trie contain the key word
+    """
+    return _retrive_branch(k, trie) is not None
+
+
+def retrie_val(k, trie):
+    key_tuple = _retrive_branch(k, trie)
     if not key_tuple:
         return None
+
     return key_tuple[1]
 
 
-def insert_key(key, v, trie_list):
-    if key == "":
-        return None
-    elif has_key(key, trie_list):
-        return None
-    else:
-        tr = trie_list
-        for char in key:
-            branch = _get_child_branch(tr, char)
-            if branch == None:
-                new_branch = [char]
-                tr.append(new_branch)
-                tr = new_branch
-            else:
-                tr = branch
-        tr.append((key, v))
-        return None
+def insert_key(key, v, trie):
+    """
+    Insert a (key, value) pair into trie
+    """
+    if not key or has_key(key, trie):
+        return
+
+    for char in key:
+        branch = _get_child_branch(trie, char)
+        if not branch:
+            new_branch = [char]
+            trie.append(new_branch)
+            trie = new_branch
+        else:
+            trie = branch
+    trie.append((key, v))
 
 
 def start_with_prefix(prefix, trie):
+    """
+    Find words start with prefix
+    """
     branch = _retrive_branch(prefix, trie)
     if not branch:
         return []
+
     prefix_list = []
     q = branch[1:]
     while q:
@@ -88,6 +103,7 @@ def start_with_prefix(prefix, trie):
             prefix_list.append(_get_bucket_key(curr_branch))
         else:
             q.extend(curr_branch[1:])
+
     return prefix_list
 
 if __name__ == "__main__":
@@ -142,7 +158,7 @@ def start_with_prefix(prefix, trie):
             Washington
             West Virginia
             Wisconsin
-            Wyoming"""    
+            Wyoming"""
     states_list = [w.strip().lower() for w in states.splitlines() if w]
     for state in states_list:
         insert_key(state, True, trie)
