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
21 def insert(self, item: Sequence[Any]):
27 >>> t.__contains__('test')
33 current = current.setdefault(child, {})
34 current[self.end] = self.end
37 def __contains__(self, item: Sequence[Any]) -> bool:
39 Check whether an item is in the Trie.
43 >>> t.__contains__('test')
45 >>> t.__contains__('testing')
51 current = self.__traverse__(item)
55 return self.end in current
57 def contains_prefix(self, item: Sequence[Any]):
59 Check whether a prefix is in the Trie. The prefix may or may not
63 >>> t.insert('testicle')
64 >>> t.contains_prefix('test')
66 >>> t.contains_prefix('testicle')
68 >>> t.contains_prefix('tessel')
72 current = self.__traverse__(item)
73 return current is not None
75 def __traverse__(self, item: Sequence[Any]):
79 current = current[child]
84 def __getitem__(self, item: Sequence[Any]):
85 """Given an item, return its Trie node which contains all
86 of the successor (child) node pointers. If the item is not
87 a node in the Trie, raise a KeyError.
91 >>> t.insert('testicle')
92 >>> t.insert('tessera')
93 >>> t.insert('tesack')
95 {'t': {'~END~': '~END~', 'i': {'c': {'l': {'e': {'~END~': '~END~'}}}}}, 's': {'e': {'r': {'a': {'~END~': '~END~'}}}}, 'a': {'c': {'k': {'~END~': '~END~'}}}}
98 ret = self.__traverse__(item)
100 raise KeyError(f"Node '{item}' is not in the trie")
103 def delete_recursively(self, node, item: Sequence[Any]) -> bool:
106 if len(node) == 0 and node is not self.root:
115 if self.delete_recursively(lower, cdr):
116 return self.delete_recursively(node, car)
119 def __delitem__(self, item: Sequence[Any]):
121 Delete an item from the Trie.
126 >>> t.insert('tessel')
130 {'t': {'e': {'s': {'t': {'~END~': '~END~'}, 's': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
131 >>> t.__delitem__('test')
135 {'t': {'e': {'s': {'s': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
140 >>> t.__delitem__('tessel')
144 {'t': {'e': {'s': {'s': {'~END~': '~END~'}}}}}
148 >>> t.__delitem__('tess')
153 >>> t.insert('testy')
159 raise KeyError(f"Node '{item}' is not in the trie")
160 self.delete_recursively(self.root, item)
165 Returns a count of the Trie's item population.
173 >>> t.insert('testicle')
181 self.content_generator = self.generate_recursively(self.root, '')
184 def generate_recursively(self, node, path: Sequence[Any]):
186 Generate items in the trie one by one.
190 >>> t.insert('tickle')
191 >>> for item in t.generate_recursively(t.root, ''):
198 if child == self.end:
201 yield from self.generate_recursively(node[child], path + child)
205 Iterate through the contents of the trie.
209 >>> t.insert('tickle')
216 ret = next(self.content_generator)
221 def successors(self, item: Sequence[Any]):
223 Return a list of the successors of an item.
229 >>> t.successors('wh')
233 >>> u.insert(['this', 'is', 'a', 'test'])
234 >>> u.insert(['this', 'is', 'a', 'robbery'])
235 >>> u.insert(['this', 'is', 'a', 'walrus'])
236 >>> u.successors(['this', 'is', 'a'])
237 ['test', 'robbery', 'walrus']
240 node = self.__traverse__(item)
243 return [x for x in node if x != self.end]
255 if node is not self.root:
256 ret = f'\n{padding}{pointer}'
266 if child != self.end:
270 if child != self.end:
277 pointer += f'{child}'
279 ret += self.repr_fancy(padding, pointer, node, node[child], has_sibling)
282 def repr_brief(self, node, delimiter):
284 A friendly string representation of the contents of the Trie.
287 >>> t.insert([10, 0, 0, 1])
288 >>> t.insert([10, 0, 0, 2])
289 >>> t.insert([10, 10, 10, 1])
290 >>> t.insert([10, 10, 10, 2])
291 >>> t.repr_brief(t.root, '.')
292 '10.[0.0.[1, 2], 10.10.[1, 2]]'
298 if child != self.end:
300 child_rep = self.repr_brief(node[child], delimiter)
301 if len(child_rep) > 0:
302 my_rep += str(child) + delimiter + child_rep + ", "
304 my_rep += str(child) + ", "
308 my_rep = f'[{my_rep}]'
313 A friendly string representation of the contents of the Trie. Under
314 the covers uses repr_fancy.
317 >>> t.insert([10, 0, 0, 1])
318 >>> t.insert([10, 0, 0, 2])
319 >>> t.insert([10, 10, 10, 1])
320 >>> t.insert([10, 10, 10, 2])
334 return self.repr_fancy('', '*', self.root, self.root, False)
337 if __name__ == '__main__':