Wikipedia Trie pseudocode in Python

Question

I translated the Wikipedia pseudo-code into Python code (little changes only were required), but I do not like the end result, as it looks like C with all those while loops and counters.

The fact is that I cannot use a for loop as i is needed to preserve state between loops.

I also added an inspect method to actually see visually what is inside the trie.

The trie built from "trasaction", "transformation", "transmission", "trance", "trap", "trip" is printed as:

     e None
    c None
           n None
          o None
         i None
        s None
       s None
      i None
     m None
          n None
         o None
        i None
       t None
      c None
     a None
             n None
            o None
           i None
          t None
         a None
        m None
       r None
      o None
     f None
    s None
   n None
   p None
  a None
   p None
  i None
 r None
t None

As you can see, the common letters at the start are written only once, and going vertically you can read the possible endings.

Let's say I want to find trance:

     e None  < -- final
    c None  < -- semifinal
           n None
          o None
         i None
        s None
       s None
      i None
     m None
          n None
         o None
        i None
       t None
      c None
     a None
             n None
            o None
           i None
          t None
         a None
        m None
       r None
      o None
     f None
    s None
   n None < -- `n` is the needed letter
   p None < -- `p` is not the needed letter, so we skip
  a None < -- `a` is the needed letter
   p None
  i None < -- `i` is not the needed letter, so we skip
 r None < -- We continue here (no choice but to)
t None < -- We start here (no choice but to)

The output is kinda reversed, but I would like to keep it so for simplicity of implementation.

import doctest

class Trie:
    """
    Implements a Trie (also known as "digital tree",  "radix tree" or "prefix tree").

    Where common starting letters of many words are stored just once
    """
    def __init__(self):
        self.child = {}

    def insert(self, s, value = 0):
        """
        Adds a given string to the trie, modifying it **in place**

        The optional value parameter is rather mysterious, my coe contains it just
        because Wikipedia's pseudocode does too.

        >>> t = Trie()
        >>> t.insert("hello")
        >>> t.insert("hi")
        >>> Trie.inspect(t)
         i None
            o None
           l None
          l None
         e None
        h None

        As you can see, the `h` was written only once, even if
        both words start with it.
        """
        node = self
        i    = 0
        n    = len(s)

        while i < n:
            try:
                _ = node.child[s[i]]
                node = node.child[s[i]]
                i = i + 1
            except KeyError:
                break

        # (* append new nodes, if necessary *)
        while i < n:
            node.child[s[i]] = Trie()
            node = node.child[s[i]]
            i = i + 1

        node.value = value

    @staticmethod
    def find(trie, word):
        """
        Returns True if the given word is the Trie, False if it is not.

        >>> t = Trie()
        >>> t.insert("example")
        >>> Trie.find(t, "example")
        True
        >>> Trie.find(t, "exemplum")
        False
        """
        if not word:
            return  True

        for char in word:

            if char == word[-1]:
                return True

            if char in trie.child:
                trie = trie.child[char]
            else:
                return False

    @staticmethod
    def inspect(trie, deepness = 0):
        """
        Shows a nicely formatted and indented Trie.
        Cannot be tested as equivalent representations are randomly chosen from.
        """
        for i in trie.child:
            print ("{}{} {}".format(
                " " * deepness, i, Trie.inspect(trie.child[i], deepness + 1)))

if __name__ == "__main__":
    doctest.testmod()

Trie on Wikipedia

(Minor note: I noticed that inspect outputs equivalent but visually different tries on different runs, if a doctest where it is used fails, do not worry.)

Answer 1

• find is not correct. If the final character occurs elsewhere in the string, the method returns True at that point. You can fix this and simplify the code at the same time:

  @staticmethod
  def find(trie, word):
      for char in word:
          if char in trie.child:
              trie = trie.child[char]
          else:
              return False
      return True
  

• The point of static methods is that they can be called without an instance of the class. Your static methods expect an instance as the first argument. Thus they could be normal instance methods instead.

• Here the first assignment is redundant. If dict lookup raises KeyError, that will occur before the assignment anyway, so the value of node is preserved even if you only use node = node.child[s[i]].

      try:
          _ = node.child[s[i]]
          node = node.child[s[i]]
          i = i + 1
      except KeyError:
          break
  

• The while loops can be changed to for loops like this for example:

  def insert(self, s, value = 0):
      node = self
      for i, char in enumerate(s):
          try:
              node = node.child[char]
          except KeyError:
              break
      else:
          i += 1
  
      # (* append new nodes, if necessary *)
      for char in s[i:]:
          node.child[char] = node = Trie()
  
      node.value = value
  

  (Note that in a chained assignment like node.child[char] = node = Trie() the right-hand-side is evaluated first and then assigned to the targets from left to right)

• You could also take advantage of collections.defaultdict and greatly simplify insert:

  def __init__(self):
      self.child = collections.defaultdict(Trie)
  
  def insert(self, s, value = 0):
      node = self
      for char in s:
          node = node.child[char]
      node.value = value
  

Answer 2

Per PEP 257:

The docstring is a phrase ending in a period. It prescribes the function or method's effect as a command ("Do this", "Return that"), not as a description; e.g. don't write "Returns the pathname ...".

So, for example, this docstring for insert:

'''Adds a given string to the trie, modifying it **in place**'''

Could be:

'''Insert the given string into the trie in place.'''

Python provides extensive operator overloading capabilities. Use them. For example, find could be called __contains__, which allows you to then check a condition like "xyx" in mytrie. You can then drop the line of the docstring about what the return value is - it is implied by using that operator. Even if you decide to leave it as a normal method, find isn't the best name - it sounds like it should return some kind of index, rather than True or False.

The reason for the value argument to insert is that a Trie behaves like a Python dict - the strings that you store efficiently are keys, and each one has an associated value. So, you could call insert instead __setitem__, and you can do things like mytrie['xyz'] = 5. If that is useful to you, you should also implement __getitem__ appropriately. If it isn't, drop that argument.

inspect is also a strange name. It sounds like it should return some kind of metadata, but instead it prints some stuff to STDOUT. Printing from support classes like this is usually worth avoiding anyway - build the full string and return it instead, and rename it to __str__ , and then to print it you just have to print(mytrie). The docstring for it seems to be lying when it says "equivalent representations are randomly chosen from" - you have no randomness here, the closest thing is that you use dicts (whose order is arbitrary but consistent, not random); it says it isn't testable, but you do test it in your doctests for another method.