3 # © Copyright 2021-2022, Scott Gasch
5 """This is a Trie class, see: https://en.wikipedia.org/wiki/Trie.
7 It attempts to follow Pythonic container patterns. See doctests
13 from typing import Any, Generator, Sequence
15 logger = logging.getLogger(__name__)
20 This is a Trie class, see: https://en.wikipedia.org/wiki/Trie.
22 It attempts to follow Pythonic container patterns. See doctests
32 self.content_generator: Generator[str] = None
34 def insert(self, item: Sequence[Any]):
40 >>> t.__contains__('test')
46 current = current.setdefault(child, {})
47 current[self.end] = self.end
50 def __contains__(self, item: Sequence[Any]) -> bool:
52 Check whether an item is in the Trie.
56 >>> t.__contains__('test')
58 >>> t.__contains__('testing')
64 current = self.__traverse__(item)
68 return self.end in current
70 def contains_prefix(self, item: Sequence[Any]):
72 Check whether a prefix is in the Trie. The prefix may or may not
76 >>> t.insert('testicle')
77 >>> t.contains_prefix('test')
79 >>> t.contains_prefix('testicle')
81 >>> t.contains_prefix('tessel')
85 current = self.__traverse__(item)
86 return current is not None
88 def __traverse__(self, item: Sequence[Any]):
92 current = current[child]
97 def __getitem__(self, item: Sequence[Any]):
98 """Given an item, return its Trie node which contains all
99 of the successor (child) node pointers. If the item is not
100 a node in the Trie, raise a KeyError.
104 >>> t.insert('testicle')
105 >>> t.insert('tessera')
106 >>> t.insert('tesack')
108 {'t': {'~END~': '~END~', 'i': {'c': {'l': {'e': {'~END~': '~END~'}}}}}, 's': {'e': {'r': {'a': {'~END~': '~END~'}}}}, 'a': {'c': {'k': {'~END~': '~END~'}}}}
111 ret = self.__traverse__(item)
113 raise KeyError(f"Node '{item}' is not in the trie")
116 def delete_recursively(self, node, item: Sequence[Any]) -> bool:
119 if len(node) == 0 and node is not self.root:
128 if self.delete_recursively(lower, cdr):
129 return self.delete_recursively(node, car)
132 def __delitem__(self, item: Sequence[Any]):
134 Delete an item from the Trie.
139 >>> t.insert('tessel')
143 {'t': {'e': {'s': {'t': {'~END~': '~END~'}, 's': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
144 >>> t.__delitem__('test')
148 {'t': {'e': {'s': {'s': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
153 >>> t.__delitem__('tessel')
157 {'t': {'e': {'s': {'s': {'~END~': '~END~'}}}}}
161 >>> t.__delitem__('tess')
166 >>> t.insert('testy')
172 raise KeyError(f"Node '{item}' is not in the trie")
173 self.delete_recursively(self.root, item)
178 Returns a count of the Trie's item population.
186 >>> t.insert('testicle')
194 self.content_generator = self.generate_recursively(self.root, '')
197 def generate_recursively(self, node, path: Sequence[Any]):
199 Generate items in the trie one by one.
203 >>> t.insert('tickle')
204 >>> for item in t.generate_recursively(t.root, ''):
211 if child == self.end:
214 yield from self.generate_recursively(node[child], path + child)
218 Iterate through the contents of the trie.
222 >>> t.insert('tickle')
229 ret = next(self.content_generator)
234 def successors(self, item: Sequence[Any]):
236 Return a list of the successors of an item.
242 >>> t.successors('wh')
246 >>> u.insert(['this', 'is', 'a', 'test'])
247 >>> u.insert(['this', 'is', 'a', 'robbery'])
248 >>> u.insert(['this', 'is', 'a', 'walrus'])
249 >>> u.successors(['this', 'is', 'a'])
250 ['test', 'robbery', 'walrus']
253 node = self.__traverse__(item)
256 return [x for x in node if x != self.end]
267 if node is not self.root:
268 ret = f'\n{padding}{pointer}'
278 if child != self.end:
282 if child != self.end:
289 pointer += f'{child}'
291 ret += self.repr_fancy(padding, pointer, node[child], has_sibling)
294 def repr_brief(self, node, delimiter):
296 A friendly string representation of the contents of the Trie.
299 >>> t.insert([10, 0, 0, 1])
300 >>> t.insert([10, 0, 0, 2])
301 >>> t.insert([10, 10, 10, 1])
302 >>> t.insert([10, 10, 10, 2])
303 >>> t.repr_brief(t.root, '.')
304 '10.[0.0.[1,2],10.10.[1,2]]'
310 if child != self.end:
312 child_rep = self.repr_brief(node[child], delimiter)
313 if len(child_rep) > 0:
314 my_rep += str(child) + delimiter + child_rep + ","
316 my_rep += str(child) + ","
320 my_rep = f'[{my_rep}]'
325 A friendly string representation of the contents of the Trie. Under
326 the covers uses repr_fancy.
329 >>> t.insert([10, 0, 0, 1])
330 >>> t.insert([10, 0, 0, 2])
331 >>> t.insert([10, 10, 10, 1])
332 >>> t.insert([10, 10, 10, 2])
346 return self.repr_fancy('', '*', self.root, False)