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
12 from typing import Any, Generator, Sequence
17 This is a Trie class, see: https://en.wikipedia.org/wiki/Trie.
19 It attempts to follow Pythonic container patterns. See doctests
29 self.content_generator: Generator[str] = None
31 def insert(self, item: Sequence[Any]):
37 >>> t.__contains__('test')
43 current = current.setdefault(child, {})
44 current[self.end] = self.end
47 def __contains__(self, item: Sequence[Any]) -> bool:
49 Check whether an item is in the Trie.
53 >>> t.__contains__('test')
55 >>> t.__contains__('testing')
61 current = self.__traverse__(item)
65 return self.end in current
67 def contains_prefix(self, item: Sequence[Any]):
69 Check whether a prefix is in the Trie. The prefix may or may not
73 >>> t.insert('testicle')
74 >>> t.contains_prefix('test')
76 >>> t.contains_prefix('testicle')
78 >>> t.contains_prefix('tessel')
82 current = self.__traverse__(item)
83 return current is not None
85 def __traverse__(self, item: Sequence[Any]):
89 current = current[child]
94 def __getitem__(self, item: Sequence[Any]):
95 """Given an item, return its Trie node which contains all
96 of the successor (child) node pointers. If the item is not
97 a node in the Trie, raise a KeyError.
101 >>> t.insert('testicle')
102 >>> t.insert('tessera')
103 >>> t.insert('tesack')
105 {'t': {'~END~': '~END~', 'i': {'c': {'l': {'e': {'~END~': '~END~'}}}}}, 's': {'e': {'r': {'a': {'~END~': '~END~'}}}}, 'a': {'c': {'k': {'~END~': '~END~'}}}}
108 ret = self.__traverse__(item)
110 raise KeyError(f"Node '{item}' is not in the trie")
113 def delete_recursively(self, node, item: Sequence[Any]) -> bool:
116 if len(node) == 0 and node is not self.root:
125 if self.delete_recursively(lower, cdr):
126 return self.delete_recursively(node, car)
129 def __delitem__(self, item: Sequence[Any]):
131 Delete an item from the Trie.
136 >>> t.insert('tessel')
140 {'t': {'e': {'s': {'t': {'~END~': '~END~'}, 's': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
141 >>> t.__delitem__('test')
145 {'t': {'e': {'s': {'s': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
150 >>> t.__delitem__('tessel')
154 {'t': {'e': {'s': {'s': {'~END~': '~END~'}}}}}
158 >>> t.__delitem__('tess')
163 >>> t.insert('testy')
169 raise KeyError(f"Node '{item}' is not in the trie")
170 self.delete_recursively(self.root, item)
175 Returns a count of the Trie's item population.
183 >>> t.insert('testicle')
191 self.content_generator = self.generate_recursively(self.root, '')
194 def generate_recursively(self, node, path: Sequence[Any]):
196 Generate items in the trie one by one.
200 >>> t.insert('tickle')
201 >>> for item in t.generate_recursively(t.root, ''):
208 if child == self.end:
211 yield from self.generate_recursively(node[child], path + child)
215 Iterate through the contents of the trie.
219 >>> t.insert('tickle')
226 ret = next(self.content_generator)
231 def successors(self, item: Sequence[Any]):
233 Return a list of the successors of an item.
239 >>> t.successors('wh')
243 >>> u.insert(['this', 'is', 'a', 'test'])
244 >>> u.insert(['this', 'is', 'a', 'robbery'])
245 >>> u.insert(['this', 'is', 'a', 'walrus'])
246 >>> u.successors(['this', 'is', 'a'])
247 ['test', 'robbery', 'walrus']
250 node = self.__traverse__(item)
253 return [x for x in node if x != self.end]
264 if node is not self.root:
265 ret = f'\n{padding}{pointer}'
275 if child != self.end:
279 if child != self.end:
286 pointer += f'{child}'
288 ret += self.repr_fancy(padding, pointer, node[child], has_sibling)
291 def repr_brief(self, node, delimiter):
293 A friendly string representation of the contents of the Trie.
296 >>> t.insert([10, 0, 0, 1])
297 >>> t.insert([10, 0, 0, 2])
298 >>> t.insert([10, 10, 10, 1])
299 >>> t.insert([10, 10, 10, 2])
300 >>> t.repr_brief(t.root, '.')
301 '10.[0.0.[1,2],10.10.[1,2]]'
307 if child != self.end:
309 child_rep = self.repr_brief(node[child], delimiter)
310 if len(child_rep) > 0:
311 my_rep += str(child) + delimiter + child_rep + ","
313 my_rep += str(child) + ","
317 my_rep = f'[{my_rep}]'
322 A friendly string representation of the contents of the Trie. Under
323 the covers uses repr_fancy.
326 >>> t.insert([10, 0, 0, 1])
327 >>> t.insert([10, 0, 0, 2])
328 >>> t.insert([10, 10, 10, 1])
329 >>> t.insert([10, 10, 10, 2])
343 return self.repr_fancy('', '*', self.root, False)
346 if __name__ == '__main__':