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.
 122
 123         Args:
 124             item: the item whose node is to be retrieved
 125
 126         Returns:
 127             A mapping that represents item in the trie.  If the
 128             keyspace of the mapping includes '~END~' a valid item
 129             ends at the node.  If the mapping contains other keys,
 130             each key indicates the presence of one or more children
 131             on the edge below the node.
 132
 133         Raises:
 134             KeyError if item is not present in the trie.
 135
 136         >>> t = Trie()
 137         >>> t.insert('test')
 138         >>> t.insert('testicle')
 139         >>> t.insert('tessera')
 140         >>> t.insert('tesack')
 141         >>> t['tes']
 142         {'t': {'~END~': '~END~', 'i': {'c': {'l': {'e': {'~END~': '~END~'}}}}}, 's': {'e': {'r': {'a': {'~END~': '~END~'}}}}, 'a': {'c': {'k': {'~END~': '~END~'}}}}
 143
 144         """
 145         ret = self.__traverse__(item)
 146         if ret is None:
 147             raise KeyError(f"Node '{item}' is not in the trie")
 148         return ret
 149
 150     def delete_recursively(self, node: Dict[Any, Any], item: Sequence[Any]) -> bool:
 151         """
 152         Deletes an item from the trie by walking the path from root to where it
 153         ends.
 154
 155         Args:
 156             node: root under which to search for item
 157             item: item whose node is the root of the recursive deletion operation
 158
 159         Returns:
 160             True if the item was not the prefix of another item such that there
 161             is nothing below item remaining anymore post delete and False if the
 162             deleted item was a proper prefix of another item in the tree such that
 163             there is still data below item remaining in the tree.
 164         """
 165         if len(item) == 1:
 166             del node[item]
 167             if len(node) == 0 and node is not self.root:
 168                 del node
 169                 return True
 170             return False
 171         else:
 172             car = item[0]
 173             cdr = item[1:]
 174             lower = node[car]
 175             ret = self.delete_recursively(lower, cdr)
 176             ret = ret and self.delete_recursively(node, car)
 177             return ret
 178
 179     def __delitem__(self, item: Sequence[Any]):
 180         """
 181         Delete an item from the Trie.
 182
 183         Args:
 184             item: the item to be deleted.
 185
 186         >>> t = Trie()
 187         >>> t.insert('test')
 188         >>> t.insert('tess')
 189         >>> t.insert('tessel')
 190         >>> len(t)
 191         3
 192         >>> t.root
 193         {'t': {'e': {'s': {'t': {'~END~': '~END~'}, 's': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
 194         >>> t.__delitem__('test')
 195         >>> len(t)
 196         2
 197         >>> t.root
 198         {'t': {'e': {'s': {'s': {'~END~': '~END~', 'e': {'l': {'~END~': '~END~'}}}}}}}
 199         >>> for x in t:
 200         ...     print(x)
 201         tess
 202         tessel
 203         >>> t.__delitem__('tessel')
 204         >>> len(t)
 205         1
 206         >>> t.root
 207         {'t': {'e': {'s': {'s': {'~END~': '~END~'}}}}}
 208         >>> for x in t:
 209         ...     print(x)
 210         tess
 211         >>> t.__delitem__('tess')
 212         >>> len(t)
 213         0
 214         >>> t.length
 215         0
 216         >>> t.root
 217         {}
 218         >>> t.insert('testy')
 219         >>> t.length
 220         1
 221         >>> len(t)
 222         1
 223         """
 224         if item not in self:
 225             raise KeyError(f"Node '{item}' is not in the trie")
 226         self.delete_recursively(self.root, item)
 227         self.length -= 1
 228
 229     def __len__(self):
 230         """
 231         Returns:
 232             A count of the trie's item population.
 233
 234         >>> t = Trie()
 235         >>> len(t)
 236         0
 237         >>> t.insert('test')
 238         >>> len(t)
 239         1
 240         >>> t.insert('tree')
 241         >>> len(t)
 242         2
 243         """
 244         return self.length
 245
 246     def __iter__(self):
 247         self.content_generator = self.generate_recursively(self.root, "")
 248         return self
 249
 250     def generate_recursively(self, node, path: Sequence[Any]):
 251         """
 252         Returns:
 253             A generator that yields the trie's items one at a time.
 254
 255         >>> t = Trie()
 256         >>> t.insert('test')
 257         >>> t.insert('tickle')
 258         >>> for item in t.generate_recursively(t.root, ''):
 259         ...     print(item)
 260         test
 261         tickle
 262
 263         """
 264         for child in node:
 265             if child == self.end:
 266                 yield path
 267             else:
 268                 yield from self.generate_recursively(node[child], path + child)
 269
 270     def __next__(self):
 271         """
 272         Iterate through the contents of the trie.
 273
 274         >>> t = Trie()
 275         >>> t.insert('test')
 276         >>> t.insert('tickle')
 277         >>> for item in t:
 278         ...     print(item)
 279         test
 280         tickle
 281
 282         """
 283         ret = next(self.content_generator)
 284         if ret is not None:
 285             return ret
 286         raise StopIteration
 287
 288     def successors(self, item: Sequence[Any]):
 289         """
 290         Returns:
 291             A list of the successors of an item.
 292
 293         >>> t = Trie()
 294         >>> t.insert('what')
 295         >>> t.insert('who')
 296         >>> t.insert('when')
 297         >>> t.successors('wh')
 298         ['a', 'o', 'e']
 299
 300         >>> u = Trie()
 301         >>> u.insert(['this', 'is', 'a', 'test'])
 302         >>> u.insert(['this', 'is', 'a', 'robbery'])
 303         >>> u.insert(['this', 'is', 'a', 'walrus'])
 304         >>> u.successors(['this', 'is', 'a'])
 305         ['test', 'robbery', 'walrus']
 306         """
 307         node = self.__traverse__(item)
 308         if node is None:
 309             return None
 310         return [x for x in node if x != self.end]
 311
 312     def _repr_fancy(
 313         self,
 314         padding: str,
 315         pointer: str,
 316         node: Any,
 317         has_sibling: bool,
 318     ) -> str:
 319         """
 320         Helper that return a fancy representation of the Trie, used by
 321         :meth:`__repr__`.
 322         """
 323         if node is None:
 324             return ""
 325         if node is not self.root:
 326             ret = f"\n{padding}{pointer}"
 327             if has_sibling:
 328                 padding += "│  "
 329             else:
 330                 padding += "   "
 331         else:
 332             ret = f"{pointer}"
 333
 334         child_count = 0
 335         for child in node:
 336             if child != self.end:
 337                 child_count += 1
 338
 339         for child in node:
 340             if child != self.end:
 341                 if child_count > 1:
 342                     pointer = "├──"
 343                     has_sibling = True
 344                 else:
 345                     pointer = "└──"
 346                     has_sibling = False
 347                 pointer += f"{child}"
 348                 child_count -= 1
 349                 ret += self._repr_fancy(padding, pointer, node[child], has_sibling)
 350         return ret
 351
 352     def repr_brief(self, node: Dict[Any, Any], delimiter: str):
 353         """
 354         A friendly string representation of the contents of the Trie.
 355
 356         Args:
 357             node: the root of the trie to represent.
 358             delimiter: character or string to stuff between edges.
 359
 360         Returns:
 361             A brief string representation of the trie.  See example.
 362
 363         >>> t = Trie()
 364         >>> t.insert([10, 0, 0, 1])
 365         >>> t.insert([10, 0, 0, 2])
 366         >>> t.insert([10, 10, 10, 1])
 367         >>> t.insert([10, 10, 10, 2])
 368         >>> t.repr_brief(t.root, '.')
 369         '10.[0.0.[1,2],10.10.[1,2]]'
 370
 371         """
 372         child_count = 0
 373         my_rep = ""
 374         for child in node:
 375             if child != self.end:
 376                 child_count += 1
 377                 child_rep = self.repr_brief(node[child], delimiter)
 378                 if len(child_rep) > 0:
 379                     my_rep += str(child) + delimiter + child_rep + ","
 380                 else:
 381                     my_rep += str(child) + ","
 382         if len(my_rep) > 1:
 383             my_rep = my_rep[:-1]
 384         if child_count > 1:
 385             my_rep = f"[{my_rep}]"
 386         return my_rep
 387
 388     def __repr__(self):
 389         """
 390         A friendly string representation of the contents of the Trie.  Under
 391         the covers uses _repr_fancy.
 392
 393         >>> t = Trie()
 394         >>> t.insert([10, 0, 0, 1])
 395         >>> t.insert([10, 0, 0, 2])
 396         >>> t.insert([10, 10, 10, 1])
 397         >>> t.insert([10, 10, 10, 2])
 398         >>> print(t)
 399         *
 400         └──10
 401            ├──0
 402            │  └──0
 403            │     ├──1
 404            │     └──2
 405            └──10
 406               └──10
 407                  ├──1
 408                  └──2
 409
 410         """
 411         return self._repr_fancy("", "*", self.root, False)
 412
 413
 414 if __name__ == "__main__":
 415     import doctest
 416
 417     doctest.testmod()