├── .gitignore ├── LICENSE ├── README ├── sample.py ├── trie.py └── trie_unittest.py /.gitignore: -------------------------------------------------------------------------------- 1 | # python stuff 2 | *.pyc 3 | 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2012 Bill Dimmick 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | -------------------------------------------------------------------------------- /README: -------------------------------------------------------------------------------- 1 | python-trie: 2 | A simple Trie implementation written in Python. 3 | 4 | - trie.py: the Python module 5 | - trie_unittest.py: unit tests for the Python module 6 | - sample.py: sample code showing various fun implementations 7 | 8 | -------------------------------------------------------------------------------- /sample.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python -B 2 | 3 | from trie import Trie 4 | 5 | """ 6 | Here's a couple of helper methods to print on assigment and dump the 7 | trie. You don't need in practice and can just treat the Trie as any 8 | other dict-like structure. 9 | """ 10 | 11 | 12 | def assign(t, k, v): 13 | print(f"Assigning {v} => t[{k}]") 14 | t[k] = v 15 | 16 | 17 | def dump(t): 18 | print("Dumping trie:") 19 | for k in t.keys(): 20 | print(f" t[{k}] => {t[k]}") 21 | 22 | 23 | """ Real Samples Start Here """ 24 | 25 | print("\nUsing a simple string as keys and numeric values...") 26 | t = Trie() 27 | assign(t, 'string1', 1) 28 | assign(t, 'string2', 2) 29 | dump(t) 30 | 31 | print("\nUsing lists as keys and bool values...") 32 | t = Trie() 33 | assign(t, [1, 2], True) 34 | assign(t, [2, 3], False) 35 | dump(t) 36 | 37 | print("\nUsing mixed types as keys and values...") 38 | t = Trie() 39 | assign(t, [1, 2], 'Hello?') 40 | assign(t, 'World', 'Earth') 41 | assign(t, 'Planet Number', 3) 42 | assign(t, (2, 1), 'X') 43 | dump(t) 44 | -------------------------------------------------------------------------------- /trie.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | class Trie(object): 4 | def __init__(self): 5 | self.path = {} 6 | self.value = None 7 | self.value_valid = False 8 | 9 | def __setitem__(self, key, value): 10 | head = key[0] 11 | if head in self.path: 12 | node = self.path[head] 13 | else: 14 | node = Trie() 15 | self.path[head] = node 16 | 17 | if len(key) > 1: 18 | remains = key[1:] 19 | node.__setitem__(remains, value) 20 | else: 21 | node.value = value 22 | node.value_valid = True 23 | 24 | def __delitem__(self, key): 25 | head = key[0] 26 | if head in self.path: 27 | node = self.path[head] 28 | if len(key) > 1: 29 | remains = key[1:] 30 | node.__delitem__(remains) 31 | else: 32 | node.value_valid = False 33 | node.value = None 34 | if len(node) == 0: 35 | del self.path[head] 36 | 37 | def __getitem__(self, key): 38 | head = key[0] 39 | if head in self.path: 40 | node = self.path[head] 41 | else: 42 | raise KeyError(key) 43 | if len(key) > 1: 44 | remains = key[1:] 45 | try: 46 | return node.__getitem__(remains) 47 | except KeyError: 48 | raise KeyError(key) 49 | elif node.value_valid: 50 | return node.value 51 | else: 52 | raise KeyError(key) 53 | 54 | def __contains__(self, key): 55 | try: 56 | self.__getitem__(key) 57 | except KeyError: 58 | return False 59 | return True 60 | 61 | def __len__(self): 62 | n = 1 if self.value_valid else 0 63 | for k in self.path.keys(): 64 | n = n + len(self.path[k]) 65 | return n 66 | 67 | def get(self, key, default=None): 68 | try: 69 | return self.__getitem__(key) 70 | except KeyError: 71 | return default 72 | 73 | def nodeCount(self): 74 | n = 0 75 | for k in self.path.keys(): 76 | n = n + 1 + self.path[k].nodeCount() 77 | return n 78 | 79 | def keys(self, prefix=[]): 80 | return self.__keys__(prefix) 81 | 82 | def __keys__(self, prefix=[], seen=[]): 83 | result = [] 84 | if self.value_valid: 85 | isStr = True 86 | val = "" 87 | for k in seen: 88 | if k is not str or len(k) > 2: 89 | isStr = False 90 | break 91 | else: 92 | val += k 93 | if isStr: 94 | result.append(val) 95 | else: 96 | result.append(prefix) 97 | if len(prefix) > 0: 98 | head = prefix[0] 99 | prefix = prefix[1:] 100 | if head in self.path: 101 | nextpaths = [head] 102 | else: 103 | nextpaths = [] 104 | else: 105 | nextpaths = self.path.keys() 106 | for k in nextpaths: 107 | nextseen = [] 108 | nextseen.extend(seen) 109 | nextseen.append(k) 110 | result.extend(self.path[k].__keys__(prefix, nextseen)) 111 | return result 112 | 113 | def __iter__(self): 114 | for k in self.keys(): 115 | yield k 116 | raise StopIteration 117 | 118 | def __add__(self, other): 119 | result = Trie() 120 | result += self 121 | result += other 122 | return result 123 | 124 | def __sub__(self, other): 125 | result = Trie() 126 | result += self 127 | result -= other 128 | return result 129 | 130 | def __iadd__(self, other): 131 | for k in other: 132 | self[k] = other[k] 133 | return self 134 | 135 | def __isub__(self, other): 136 | for k in other: 137 | del self[k] 138 | return self 139 | -------------------------------------------------------------------------------- /trie_unittest.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python -B 2 | 3 | import unittest 4 | from trie import Trie 5 | 6 | 7 | class TestTrie(unittest.TestCase): 8 | def setUp(self): 9 | self.trie = Trie() 10 | 11 | def _square_brackets(self, key): 12 | return self.trie[key] 13 | 14 | def test_basicAssignment(self): 15 | self.trie["Foo"] = True 16 | self.assertTrue(self.trie["Foo"]) 17 | self.assertRaises(KeyError, self._square_brackets, "Food") 18 | self.assertEquals(1, len(self.trie)) 19 | self.assertEquals(3, self.trie.nodeCount()) 20 | self.assertTrue("Foo" in self.trie) 21 | self.trie["Bar"] = None 22 | self.assertTrue("Bar" in self.trie) 23 | 24 | def test_basicRemoval(self): 25 | self.trie["Foo"] = True 26 | self.assertTrue(self.trie["Foo"]) 27 | del self.trie["Foo"] 28 | self.assertRaises(KeyError, self._square_brackets, "Foo") 29 | self.assertEquals(0, len(self.trie)) 30 | self.assertEquals(0, self.trie.nodeCount()) 31 | self.assertFalse("Foo" in self.trie) 32 | 33 | def test_MixedTypes(self): 34 | self.trie["Foo"] = True 35 | self.trie[[1, 2, 3]] = True 36 | self.assertTrue(self.trie["Foo"]) 37 | self.assertTrue(self.trie[[1, 2, 3]]) 38 | self.assertTrue([1, 2, 3] in self.trie) 39 | self.assertTrue("Foo" in self.trie) 40 | del self.trie[[1, 2, 3]] 41 | self.assertFalse([1, 2, 3] in self.trie) 42 | 43 | def test_Iteration(self): 44 | self.trie["Foo"] = True 45 | self.trie["Bar"] = True 46 | self.trie["Grok"] = True 47 | for k in self.trie: 48 | self.assertTrue(k in self.trie) 49 | self.assertTrue(self.trie[k]) 50 | 51 | def test_Addition(self): 52 | self.trie["Foo"] = True 53 | t2 = Trie() 54 | t2["Food"] = True 55 | t3 = t2 + self.trie 56 | self.assertTrue("Foo" in self.trie) 57 | self.assertFalse("Food" in self.trie) 58 | self.assertTrue("Food" in t2) 59 | self.assertFalse("Foo" in t2) 60 | self.assertTrue("Foo" in t3) 61 | self.assertTrue("Food" in t3) 62 | 63 | def test_Subtraction(self): 64 | self.trie["Food"] = True 65 | self.trie["Foo"] = True 66 | t2 = Trie() 67 | t2["Food"] = True 68 | t3 = self.trie - t2 69 | t4 = t2 - self.trie 70 | self.assertTrue("Food" in self.trie) 71 | self.assertTrue("Foo" in self.trie) 72 | self.assertTrue("Food" in t2) 73 | self.assertTrue("Foo" in t3) 74 | self.assertFalse("Food" in t3) 75 | self.assertFalse("Foo" in t4) 76 | self.assertFalse("Food" in t4) 77 | 78 | def test_SelfAdd(self): 79 | self.trie["Foo"] = True 80 | t2 = Trie() 81 | t2["Food"] = True 82 | self.assertTrue("Foo" in self.trie) 83 | self.assertFalse("Food" in self.trie) 84 | self.assertTrue("Food" in t2) 85 | self.assertFalse("Foo" in t2) 86 | self.trie += t2 87 | self.assertTrue("Foo" in self.trie) 88 | self.assertTrue("Food" in self.trie) 89 | 90 | def test_SelfSub(self): 91 | self.trie["Foo"] = True 92 | self.trie["Food"] = True 93 | t2 = Trie() 94 | t2["Food"] = True 95 | self.assertTrue("Food" in self.trie) 96 | self.assertTrue("Foo" in self.trie) 97 | self.assertTrue("Food" in t2) 98 | self.trie -= t2 99 | self.assertFalse("Food" in self.trie) 100 | self.assertTrue("Foo" in self.trie) 101 | self.assertTrue("Food" in t2) 102 | 103 | def test_SelfGet(self): 104 | self.trie["Foo"] = True 105 | self.assertTrue(self.trie["Foo"]) 106 | self.assertRaises(KeyError, self._square_brackets, "Food") 107 | self.assertEquals("Bar", self.trie.get("Food", "Bar")) 108 | self.assertEquals("Bar", self.trie.get("Food", default="Bar")) 109 | self.assertTrue(self.trie.get("Foo")) 110 | self.assertTrue(self.trie.get("Food") is None) 111 | 112 | def test_KeysByPrefix(self): 113 | self.trie["Foo"] = True 114 | self.trie["Food"] = True 115 | self.trie["Eggs"] = True 116 | kset = self.trie.keys() 117 | self.assertTrue("Foo" in kset) 118 | self.assertTrue("Food" in kset) 119 | self.assertTrue("Eggs" in kset) 120 | kset = self.trie.keys("Foo") 121 | self.assertTrue("Foo" in kset) 122 | self.assertTrue("Food" in kset) 123 | kset = self.trie.keys("Ox") 124 | self.assertEquals(0, len(kset)) 125 | 126 | 127 | if __name__ == '__main__': 128 | unittest.main() 129 | --------------------------------------------------------------------------------