3 from typing import Any, Sequence
8 This is a Trie class, see: https://en.wikipedia.org/wiki/Trie.
10 It attempts to follow Pythonic container patterns. See doctests
20 def insert(self, item: Sequence[Any]):
26 >>> t.__contains__('test')
32 current = current.setdefault(child, {})
33 current[self.end] = self.end
36 def __contains__(self, item: Sequence[Any]) -> bool:
38 Check whether an item is in the Trie.
42 >>> t.__contains__('test')
44 >>> t.__contains__('testing')
50 current = self.__traverse__(item)
54 return self.end in current
56 def contains_prefix(self, item: Sequence[Any]):
58 Check whether a prefix is in the Trie. The prefix may or may not
62 >>> t.insert('testicle')
63 >>> t.contains_prefix('test')
65 >>> t.contains_prefix('testicle')
67 >>> t.contains_prefix('tessel')
71 current = self.__traverse__(item)
72 return current is not None
74 def __traverse__(self, item: Sequence[Any]):
78 current = current[child]
83 def __getitem__(self, item: Sequence[Any]):
84 """Given an item, return its Trie node which contains all
85 of the successor (child) node pointers. If the item is not
86 a node in the Trie, raise a KeyError.
90 >>> t.insert('testicle')
91 >>> t.insert('tessera')
92 >>> t.insert('tesack')
94 {'t': {'~END~': '~END~', 'i': {'c': {'l': {'e': {'~END~': '~END~'}}}}}, 's': {'e': {'r': {'a': {'~END~': '~END~'}}}}, 'a': {'c': {'k': {'~END~': '~END~'}}}}
97 ret = self.__traverse__(item)
99 raise KeyError(f"Node '{item}' is not in the trie")
102 def delete_recursively(self, node, item: Sequence[Any]) -> bool:
105 if len(node) == 0 and node is not self.root:
114 if self.delete_recursively(lower, cdr):
115 return self.delete_recursively(node, car)
118 def __delitem__(self, item: Sequence[Any]):
120 Delete an item from the Trie.
125 >>> t.insert('tessel')
129 {'t': {'e': {'s': {'t': {'~END~': '~END~'}, 's': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
130 >>> t.__delitem__('test')
134 {'t': {'e': {'s': {'s': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
139 >>> t.__delitem__('tessel')
143 {'t': {'e': {'s': {'s': {'~END~': '~END~'}}}}}
147 >>> t.__delitem__('tess')
152 >>> t.insert('testy')
158 raise KeyError(f"Node '{item}' is not in the trie")
159 self.delete_recursively(self.root, item)
164 Returns a count of the Trie's item population.
172 >>> t.insert('testicle')
180 self.content_generator = self.generate_recursively(self.root, '')
183 def generate_recursively(self, node, path: Sequence[Any]):
185 Generate items in the trie one by one.
189 >>> t.insert('tickle')
190 >>> for item in t.generate_recursively(t.root, ''):
197 if child == self.end:
200 yield from self.generate_recursively(node[child], path + child)
204 Iterate through the contents of the trie.
208 >>> t.insert('tickle')
215 ret = next(self.content_generator)
220 def successors(self, item: Sequence[Any]):
222 Return a list of the successors of an item.
228 >>> t.successors('wh')
232 >>> u.insert(['this', 'is', 'a', 'test'])
233 >>> u.insert(['this', 'is', 'a', 'robbery'])
234 >>> u.insert(['this', 'is', 'a', 'walrus'])
235 >>> u.successors(['this', 'is', 'a'])
236 ['test', 'robbery', 'walrus']
239 node = self.__traverse__(item)
242 return [x for x in node if x != self.end]
244 def repr_fancy(self, padding: str, pointer: str, parent: str, node: Any, has_sibling: bool):
247 if node is not self.root:
248 ret = f'\n{padding}{pointer}'
258 if child != self.end:
262 if child != self.end:
269 pointer += f'{child}'
271 ret += self.repr_fancy(padding, pointer, node, node[child], has_sibling)
274 def repr_brief(self, node, delimiter):
276 A friendly string representation of the contents of the Trie.
279 >>> t.insert([10, 0, 0, 1])
280 >>> t.insert([10, 0, 0, 2])
281 >>> t.insert([10, 10, 10, 1])
282 >>> t.insert([10, 10, 10, 2])
283 >>> t.repr_brief(t.root, '.')
284 '10.[0.0.[1, 2], 10.10.[1, 2]]'
290 if child != self.end:
292 child_rep = self.repr_brief(node[child], delimiter)
293 if len(child_rep) > 0:
294 my_rep += str(child) + delimiter + child_rep + ", "
296 my_rep += str(child) + ", "
300 my_rep = f'[{my_rep}]'
305 A friendly string representation of the contents of the Trie. Under
306 the covers uses repr_fancy.
309 >>> t.insert([10, 0, 0, 1])
310 >>> t.insert([10, 0, 0, 2])
311 >>> t.insert([10, 10, 10, 1])
312 >>> t.insert([10, 10, 10, 2])
326 return self.repr_fancy('', '*', self.root, self.root, False)
329 if __name__ == '__main__':