src/pyutils/collectionz/trie.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2023, Scott Gasch
   4
   5 """This module contains the implementation of a Trie tree (or prefix
   6 tree).  See: https://en.wikipedia.org/wiki/Trie.
   7
   8 It can be used with arbitrary sequences as keys and stores its values
   9 in a tree with paths determined by the sequence determined by each
  10 keys' sequences.  Thus, it can determine whether a given value is
  11 contained in the tree via a simple traversal in :math:`O(n)` where n
  12 is the number of steps in the item's sequence and can also check
  13 whether a key-prefix is present in the tree in :math:`O(n)` time.
  14
  15 Given a node in the BST, it is easy to determine all items that are
  16 stored beneath that node.  See examples below.
  17
  18 """
  19
  20 import logging
  21 from typing import Any, Dict, Generator, Sequence
  22
  23 logger = logging.getLogger(__name__)
  24
  25
  26 class Trie(object):
  27     """
  28     This is a Trie class, see: https://en.wikipedia.org/wiki/Trie.
  29
  30     It attempts to follow Pythonic container patterns.  See doctests
  31     for examples.
  32     """
  33
  34     def __init__(self):
  35         """Create an empty trie."""
  36         self.root = {}
  37         self.end = "~END~"
  38         self.length = 0
  39         self.viz = ""
  40         self.content_generator: Generator[str] = None
  41
  42     def insert(self, item: Sequence[Any]) -> None:
  43         """
  44         Insert an item into the trie.  Items are represented by a :class:`Sequence`
  45         where each item in the sequence describes a set in the path to the item.
  46
  47         Args:
  48             item: the item to be inserted.
  49
  50         >>> t = Trie()
  51         >>> t.insert('test')
  52         >>> t.__contains__('test')
  53         True
  54         """
  55         current = self.root
  56         for child in item:
  57             current = current.setdefault(child, {})
  58         current[self.end] = self.end
  59         self.length += 1
  60
  61     def __contains__(self, item: Sequence[Any]) -> bool:
  62         """
  63         Checks whether an item is in the Trie.
  64
  65         Args:
  66             item: the item whose presence is to be determined.
  67
  68         Returns:
  69             True if `item` is present, False otherwise.
  70
  71         >>> t = Trie()
  72         >>> t.insert('test')
  73         >>> t.__contains__('test')
  74         True
  75         >>> t.__contains__('testing')
  76         False
  77         >>> 'test' in t
  78         True
  79         """
  80         current = self.__traverse__(item)
  81         if current is None:
  82             return False
  83         else:
  84             return self.end in current
  85
  86     def contains_prefix(self, item: Sequence[Any]):
  87         """
  88         Check whether a prefix is in the Trie.  The prefix may or may not
  89         be a full item.
  90
  91         Args:
  92             item: the item describing the prefix to be checked.
  93
  94         Returns:
  95             True if the prefix described by `item` is present, False otherwise.
  96
  97         >>> t = Trie()
  98         >>> t.insert('tricycle')
  99         >>> t.contains_prefix('tri')
 100         True
 101         >>> t.contains_prefix('tricycle')
 102         True
 103         >>> t.contains_prefix('triad')
 104         False
 105
 106         """
 107         current = self.__traverse__(item)
 108         return current is not None
 109
 110     def __traverse__(self, item: Sequence[Any]):
 111         current = self.root
 112         for child in item:
 113             if child in current:
 114                 current = current[child]
 115             else:
 116                 return None
 117         return current
 118
 119     def __getitem__(self, item: Sequence[Any]) -> Dict[Any, Any]:
 120         """Given an item, return its trie node which contains all
 121         of the successor (child) node pointers.  If the item is not
 122         a node in the Trie, raise a KeyError.
 123
 124         Args:
 125             item: the item whose node is to be retrieved
 126
 127         Returns:
 128             A mapping that represents item in the trie.  If the
 129             keyspace of the mapping includes '~END~' a valid item
 130             ends at the node.  If the mapping contains other keys,
 131             each key indicates the presence of one or more children
 132             on the edge below the node.
 133
 134         Raises:
 135             KeyError if item is not present in the trie.
 136
 137         >>> t = Trie()
 138         >>> t.insert('test')
 139         >>> t.insert('testicle')
 140         >>> t.insert('tessera')
 141         >>> t.insert('tesack')
 142         >>> t['tes']
 143         {'t': {'~END~': '~END~', 'i': {'c': {'l': {'e': {'~END~': '~END~'}}}}}, 's': {'e': {'r': {'a': {'~END~': '~END~'}}}}, 'a': {'c': {'k': {'~END~': '~END~'}}}}
 144
 145         """
 146         ret = self.__traverse__(item)
 147         if ret is None:
 148             raise KeyError(f"Node '{item}' is not in the trie")
 149         return ret
 150
 151     def delete_recursively(self, node: Dict[Any, Any], item: Sequence[Any]) -> bool:
 152         """
 153         Deletes an item from the trie by walking the path from root to where it
 154         ends.
 155
 156         Args:
 157             node: root under which to search for item
 158             item: item whose node is the root of the recursive deletion operation
 159
 160         Returns:
 161             True if the item was not the prefix of another item such that there
 162             is nothing below item remaining anymore post delete and False if the
 163             deleted item was a proper prefix of another item in the tree such that
 164             there is still data below item remaining in the tree.
 165         """
 166         if len(item) == 1:
 167             del node[item]
 168             if len(node) == 0 and node is not self.root:
 169                 del node
 170                 return True
 171             return False
 172         else:
 173             car = item[0]
 174             cdr = item[1:]
 175             lower = node[car]
 176             ret = self.delete_recursively(lower, cdr)
 177             ret = ret and self.delete_recursively(node, car)
 178             return ret
 179
 180     def __delitem__(self, item: Sequence[Any]):
 181         """
 182         Delete an item from the Trie.
 183
 184         Args:
 185             item: the item to be deleted.
 186
 187         >>> t = Trie()
 188         >>> t.insert('test')
 189         >>> t.insert('tess')
 190         >>> t.insert('tessel')
 191         >>> len(t)
 192         3
 193         >>> t.root
 194         {'t': {'e': {'s': {'t': {'~END~': '~END~'}, 's': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
 195         >>> t.__delitem__('test')
 196         >>> len(t)
 197         2
 198         >>> t.root
 199         {'t': {'e': {'s': {'s': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
 200         >>> for x in t:
 201         ...     print(x)
 202         tess
 203         tessel
 204         >>> t.__delitem__('tessel')
 205         >>> len(t)
 206         1
 207         >>> t.root
 208         {'t': {'e': {'s': {'s': {'~END~': '~END~'}}}}}
 209         >>> for x in t:
 210         ...     print(x)
 211         tess
 212         >>> t.__delitem__('tess')
 213         >>> len(t)
 214         0
 215         >>> t.length
 216         0
 217         >>> t.root
 218         {}
 219         >>> t.insert('testy')
 220         >>> t.length
 221         1
 222         >>> len(t)
 223         1
 224         """
 225         if item not in self:
 226             raise KeyError(f"Node '{item}' is not in the trie")
 227         self.delete_recursively(self.root, item)
 228         self.length -= 1
 229
 230     def __len__(self):
 231         """
 232         Returns:
 233             A count of the trie's item population.
 234
 235         >>> t = Trie()
 236         >>> len(t)
 237         0
 238         >>> t.insert('test')
 239         >>> len(t)
 240         1
 241         >>> t.insert('tree')
 242         >>> len(t)
 243         2
 244         """
 245         return self.length
 246
 247     def __iter__(self):
 248         self.content_generator = self.generate_recursively(self.root, "")
 249         return self
 250
 251     def generate_recursively(self, node, path: Sequence[Any]):
 252         """
 253         Returns:
 254             A generator that yields the trie's items one at a time.
 255
 256         >>> t = Trie()
 257         >>> t.insert('test')
 258         >>> t.insert('tickle')
 259         >>> for item in t.generate_recursively(t.root, ''):
 260         ...     print(item)
 261         test
 262         tickle
 263
 264         """
 265         for child in node:
 266             if child == self.end:
 267                 yield path
 268             else:
 269                 yield from self.generate_recursively(node[child], path + child)
 270
 271     def __next__(self):
 272         """
 273         Iterate through the contents of the trie.
 274
 275         >>> t = Trie()
 276         >>> t.insert('test')
 277         >>> t.insert('tickle')
 278         >>> for item in t:
 279         ...     print(item)
 280         test
 281         tickle
 282
 283         """
 284         ret = next(self.content_generator)
 285         if ret is not None:
 286             return ret
 287         raise StopIteration
 288
 289     def successors(self, item: Sequence[Any]):
 290         """
 291         Returns:
 292             A list of the successors of an item.
 293
 294         >>> t = Trie()
 295         >>> t.insert('what')
 296         >>> t.insert('who')
 297         >>> t.insert('when')
 298         >>> t.successors('wh')
 299         ['a', 'o', 'e']
 300
 301         >>> u = Trie()
 302         >>> u.insert(['this', 'is', 'a', 'test'])
 303         >>> u.insert(['this', 'is', 'a', 'robbery'])
 304         >>> u.insert(['this', 'is', 'a', 'walrus'])
 305         >>> u.successors(['this', 'is', 'a'])
 306         ['test', 'robbery', 'walrus']
 307         """
 308         node = self.__traverse__(item)
 309         if node is None:
 310             return None
 311         return [x for x in node if x != self.end]
 312
 313     def _repr_fancy(
 314         self,
 315         padding: str,
 316         pointer: str,
 317         node: Any,
 318         has_sibling: bool,
 319     ) -> str:
 320         """
 321         Helper that return a fancy representation of the Trie, used by
 322         :meth:`__repr__`.
 323         """
 324         if node is None:
 325             return ""
 326         if node is not self.root:
 327             ret = f"\n{padding}{pointer}"
 328             if has_sibling:
 329                 padding += "│  "
 330             else:
 331                 padding += "   "
 332         else:
 333             ret = f"{pointer}"
 334
 335         child_count = 0
 336         for child in node:
 337             if child != self.end:
 338                 child_count += 1
 339
 340         for child in node:
 341             if child != self.end:
 342                 if child_count > 1:
 343                     pointer = "├──"
 344                     has_sibling = True
 345                 else:
 346                     pointer = "└──"
 347                     has_sibling = False
 348                 pointer += f"{child}"
 349                 child_count -= 1
 350                 ret += self._repr_fancy(padding, pointer, node[child], has_sibling)
 351         return ret
 352
 353     def repr_brief(self, node: Dict[Any, Any], delimiter: str):
 354         """
 355         A friendly string representation of the contents of the Trie.
 356
 357         Args:
 358             node: the root of the trie to represent.
 359             delimiter: character or string to stuff between edges.
 360
 361         Returns:
 362             A brief string representation of the trie.  See example.
 363
 364         >>> t = Trie()
 365         >>> t.insert([10, 0, 0, 1])
 366         >>> t.insert([10, 0, 0, 2])
 367         >>> t.insert([10, 10, 10, 1])
 368         >>> t.insert([10, 10, 10, 2])
 369         >>> t.repr_brief(t.root, '.')
 370         '10.[0.0.[1,2],10.10.[1,2]]'
 371
 372         """
 373         child_count = 0
 374         my_rep = ""
 375         for child in node:
 376             if child != self.end:
 377                 child_count += 1
 378                 child_rep = self.repr_brief(node[child], delimiter)
 379                 if len(child_rep) > 0:
 380                     my_rep += str(child) + delimiter + child_rep + ","
 381                 else:
 382                     my_rep += str(child) + ","
 383         if len(my_rep) > 1:
 384             my_rep = my_rep[:-1]
 385         if child_count > 1:
 386             my_rep = f"[{my_rep}]"
 387         return my_rep
 388
 389     def __repr__(self):
 390         """
 391         A friendly string representation of the contents of the Trie.  Under
 392         the covers uses _repr_fancy.
 393
 394         >>> t = Trie()
 395         >>> t.insert([10, 0, 0, 1])
 396         >>> t.insert([10, 0, 0, 2])
 397         >>> t.insert([10, 10, 10, 1])
 398         >>> t.insert([10, 10, 10, 2])
 399         >>> print(t)
 400         *
 401         └──10
 402            ├──0
 403            │  └──0
 404            │     ├──1
 405            │     └──2
 406            └──10
 407               └──10
 408                  ├──1
 409                  └──2
 410
 411         """
 412         return self._repr_fancy("", "*", self.root, False)
 413
 414
 415 if __name__ == "__main__":
 416     import doctest
 417
 418     doctest.testmod()