Readability and Performance of my Trie implementation

Question

As an interesting exercise in Python 3.3, I implemented a Trie (also known as a Prefix Tree).

Example usages

Create from list of key-value-tuples

mapping = (('she', 1), ('sells', 5), ('sea', 10), ('shells', 19), ('today', 5))
trie = Trie.from_mapping(mapping)
some_prefix = 'sh'
some_key = 'today'
print('Keys:', *trie)
print('Items: ', *trie.items())
print('Keys prefixed "{}":'.format(some_prefix), *trie.keys(prefix=some_prefix))
trie[some_key] = -55
print('Value of "{}":'.format(some_key), trie[some_key])
print(some_key, 'is in' if some_key in trie else 'is not in', 'the Trie')

Output

Keys: today sells sea she shells
Items: ('today', 5) ('sells', 5) ('sea', 10) ('she', 1) ('shells', 19)
Keys prefixed "sh": she shells
Value of "today": -55
today is in the Trie

Find out lengths of words in a string

words = 'She sells sea shells today she sells no shells she sells'.split()
trie = Trie.from_keys(words, value_selector=len)
print(*trie.items())

Output

('sells', 5) ('sea', 3) ('she', 3) ('shells', 6) ('today', 5) ('no', 2) ('She', 3)

---

Implementation

For simplicity, each entry is represented by a dict mapping characters to child entries – if this is problematic, feel free to remark on that.

class _Entry:
 def __init__(self, value=None):
 self.value = value
 self.children = {}
 def __iter__(self):
 return iter(self.children.items())

I made it iterable so I don't have to constantly write things like entry.children.items().

Here's the Trie itself – please feel free to point out any or all issues you can spot. I started taking a look at Python two weeks ago, so I'd also appreciate your comments on how Pythonic (or not) this is, as well as any potential issues of performance.

class Trie:
 @staticmethod
 def from_mapping(mapping):
 """
 Constructs a Trie from an iterable of two-item tuples.
 :param mapping: An iterable of key value pairs (each key must itself
 be an iterable)
 :return: A new Trie containing the keys and values in the mapping
 """
 trie = Trie()
 for key, value in mapping:
 if not key:
 continue
 trie[key] = value
 return trie
 @staticmethod
 def from_keys(keys, value_selector):
 """
 Constructs a Trie from an iterable of keys, mapping each key to the
 value returned by the value selector function, which is passed each key.
 :param keys: an iterable of keys (each key must itself be an iterable)
 :return: a new Trie containing the keys in the mapping
 """
 trie = Trie()
 for key in keys:
 if not key:
 continue
 trie[key] = value_selector(key)
 return trie
 def __init__(self):
 """
 Creates an empty Trie.
 """
 self._root = _Entry()
 self._length = 0
 def __getitem__(self, key):
 """
 Gets the value corresponding to the specified key.
 :param key: an iterable
 :return: the value corresponding to the key if it is in the Trie,
 otherwise None
 """
 entry = self._find(key)
 if entry is not None:
 return entry.value
 def __setitem__(self, key, value):
 """
 Sets or inserts the specified value, mapping the specified key to it
 :param key: an iterable to update or insert
 :param value: the object to associate with the key
 """
 self.update(key, lambda old_value: value)
 def __delitem__(self, key):
 """
 Removes the value of the specified key without removing it from the Trie
 :param key: the iterable whose value is to be removed
 """
 self[key] = None
 self._length -= 1
 def __contains__(self, key):
 """
 Determines if the key is in the Trie and has a value associated with it
 :param key: the iterable to search for
 :return: True if the Trie contains the key, otherwise False
 """
 entry = self._find(key)
 return entry is not None and entry.value is not None
 def __len__(self):
 return self._length
 def __bool__(self):
 return len(self) > 0
 def __iter__(self):
 return iter(self.keys())
 def keys(self, prefix=''):
 """
 Searches for keys beginning with the specified prefix
 :param prefix: an iterable
 :return: A generator yielding the keys one by one, as they are found
 """
 yield from (key for key, value in self.items(prefix))
 def items(self, prefix=''):
 """
 Searches for key value pairs where the keys begin with the specified
 prefix
 :param prefix: an iterable
 :return: A generator yielding tuples of keys and values as they are
 found
 """
 start = self._find(prefix)
 if start is None:
 return ()
 stack = [(prefix, start)]
 while stack:
 current_prefix, entry = stack.pop()
 if entry.value is not None:
 yield current_prefix, entry.value
 for char, children in entry:
 stack.append((current_prefix + char, children))
 def update(self, key, value_selector):
 """
 Updates key with the value obtained from the value selector function,
 to which the old value corresponding to the key is passed.
 :param key: the key whose value is to be updated
 :param value_selector: a function that takes the old value as an
 argument and returns a new value
 """
 entry = self._root
 for char in key:
 if char not in entry.children:
 entry.children[char] = _Entry()
 entry = entry.children[char]
 new_value = value_selector(entry.value)
 if entry.value != new_value:
 entry.value = new_value
 self._length += 1
 def _find(self, key):
 entry = self._root
 for char in key:
 if char not in entry.children:
 return None
 entry = entry.children[char]
 return entry

Answer 1

1. Bug

1. Sometimes the _length attribute is updated wrongly. First, when deleting an item, the length is decremented regardless of whether the item was actually present in the trie:

   >>> trie = Trie()
   >>> del trie['key']
   >>> len(trie)
   Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
   ValueError: __len__() should return >= 0
   

   Second, when updating an item, the length is incremented regardless of whether the item was actually new:

   >>> trie = Trie()
   >>> trie['key'] = 1
   >>> trie['key'] = 2
   >>> print(*trie)
   key
   >>> len(trie)
   2
   

2. Other comments

1. You have docstrings for all your methods: awesome! You might consider adding a docstring for the Trie class (look at the output of help(Trie) and consider what else is needed). This is also a good class to consider adding doctests.

2. Since the class _Entry is only used in Trie, you could make it an attribute of Trie.

3. In the _Entry class, you use None as a placeholder to mean that there is no value associated with this entry. This prevents anyone from storing None as a value in a Trie. Instead, you could leave the value attribute unset to indicate no value (this can be detected with the hasattr function), and write del entry.value to delete it.

4. You could have _Entry inherit from dict and then you wouldn't need to have an attribute children.

5. By using @staticmethod you prevent anyone from subclassing Trie, because when someone calls SubTrie.from_mapping, the object they get back will be a Trie rather than a SubTrie. You should use @classmethod instead:

   @classmethod
   def from_mapping(class_, mapping):
    trie = class_()
    # ...
   

6. It's not clear why you have the lines:

   if not key:
    continue
   

   in from_mapping and from_keys. Why can't I add the empty string to a Trie? (If you have a good reason, it would be worth mentioning it in the documentation, and also worth raising an error here, instead of silently discarding these keys.)

7. The __getitem__ method returns None to indicate failure. In Python it's normal to raise an exception to indicate failure rather than returning an exceptional type. In this case it would be appropriate to raise KeyError (just as the built-in dict.__getitem__ does).

8. You might consider whether Trie should inherit from MutableMapping in the collections.abc module: this would give you values, get, __eq__, __ne__, pop, popitem, clear, update and setdefault methods for free.

9. Since your Trie class presents a mapping-like interface, Python programmers will expect its update method to present a similar interface to the update method on ordinary dicationary objects. (This is easy if you inherit from MutableMapping since the update method is defined in that abstract base class.)

   (Your original update would then need a new name.)

10. Similarly, since your Trie class presents a mapping-like interface, Python programmers will expect its constructor to present a similar interface to the dict constructor. (Again, this is easy if you inherit from MutableMapping since you can just pass the positional and keyword arguments to the update method.)

11. It seems perverse that the value_selector argument to the update method is different from the value_selector argument to the from_keys method. (The former takes the old value as its argument but the latter takes the key as the argument.) This can only be confusing to users.

    Also, if your update method is called, but there's no value in the trie for key, the value_selector function gets called anyway (with the argument None). You might consider whether this method should take a third argument, the value to use in the case where there is no value in the trie. (Or maybe you should raise KeyError in this case.)

    I have a feeling that your from_keys and update methods are not very Pythonic. Instead of from_keys, callers would be happy to write:

    Trie((key, f(key)) for key in keys)
    

    and instead of trie.update(key, lambda old_value: old_value + 1) caller would be happy to write

    trie[key] += 1
    

    My guess is that your goal in writing the update method is to avoid the double lookup in the above. So it's a reasonable thing to provide, it just needs some thought about the name and interface.

12. You could simplify some of your logic if your _find method was able to create new _Entry objects. See below for how.

13. If your function body consists only of yield from <generator>, you might as well write return <generator>.

3. Revised code

This is how I would write the class, addressing all the points above.

from collections.abc import MutableMapping
class Trie(MutableMapping):
 """
 A prefix tree.
 >>> mapping = dict(she=1, sells=5, sea=10, shells=19, today=5)
 >>> trie = Trie(mapping)
 >>> some_prefix = 'sh'
 >>> some_key = 'today'
 >>> print(*sorted(trie))
 sea sells she shells today
 >>> print(*sorted(trie.items()))
 ('sea', 10) ('sells', 5) ('she', 1) ('shells', 19) ('today', 5)
 >>> print(*trie.keys(prefix='sh'))
 she shells
 >>> trie['today'] = -55
 >>> trie['today']
 -55
 >>> len(trie)
 5
 >>> 'shells' in trie, 'shore' in trie
 (True, False)
 >>> del trie['shells']
 >>> len(trie)
 4
 """
 # A dictionary with an optional :value: attribute.
 class _Entry(dict):
 def has_value(self):
 return hasattr(self, 'value')
 def __init__(self, *args, **kwargs):
 """
 Return a new True initialized from an optional positional
 argument and a possibly empty set of keyword arguments, as for
 the built-in type :dict:
 """
 self._root = self._Entry()
 self._length = 0
 self.update(*args, **kwargs)
 def __getitem__(self, key):
 """
 Gets the value corresponding to the specified key.
 :param key: an iterable
 :return: the value corresponding to the key if it is in the Trie,
 otherwise raise KeyError.
 """
 entry = self._find(key)
 if entry.has_value():
 return entry.value
 else:
 raise KeyError(key)
 def __setitem__(self, key, value):
 """
 Sets or inserts the specified value, mapping the specified key to it
 :param key: an iterable to update or insert
 :param value: the object to associate with the key
 """
 entry = self._find(key, create=True)
 if not entry.has_value():
 self._length += 1
 entry.value = value
 def __delitem__(self, key):
 """
 Removes the value of the specified key from the Trie
 :param key: the iterable whose value is to be removed
 """
 entry = self._find(key)
 if entry.has_value():
 del entry.value
 self._length -= 1
 else:
 raise KeyError(key)
 def __len__(self):
 return self._length
 def __iter__(self):
 return iter(self.keys())
 def keys(self, prefix=''):
 """
 Searches for keys beginning with the specified prefix
 :param prefix: an iterable
 :return: A generator yielding the keys one by one, as they are found
 """
 return (key for key, value in self.items(prefix))
 def items(self, prefix=''):
 """
 Searches for key value pairs where the keys begin with the specified
 prefix
 :param prefix: an iterable
 :return: A generator yielding tuples of keys and values as they are
 found
 """
 try:
 start = self._find(prefix)
 except KeyError:
 raise StopIteration
 stack = [(prefix, start)]
 while stack:
 current_prefix, entry = stack.pop()
 if entry.has_value():
 yield current_prefix, entry.value
 for char, children in entry.items():
 stack.append((current_prefix + char, children))
 def _find(self, key, create=False):
 entry = self._root
 for char in key:
 if char in entry:
 entry = entry[char]
 elif create:
 new_entry = self._Entry()
 entry[char] = new_entry
 entry = new_entry
 else:
 raise KeyError(key)
 return entry

4. Algorithmic concerns

1. You create _Entry objects for all elements in each key, even if there are no other keys sharing any of those entries. For example, if I create an empty Trie and then add a value for the key supercalifragilisticexpialidocious, 34 _Entry objects will be constructed, and when I look up this key, all 34 of them will have to be traversed. You might consider storing the tail of a key differently when it doesn't need to share space with other keys, and creating _Entry objects dynamically as necessary when you add new keys that share a prefix with existing keys.

2. When you delete keys, you don't remove _Entry objects that have become empty as a result. This means that later lookups may need to descend deeper into the trie before discovering that the key is not present.

• Winston Ewert: Awesome! My only issue is that I sternly disapprove of the suggesting to subclass dict for the _Entry class.
• Gareth Rees: @WinstonEwert: Can you explain your concern?
• Adam: Gareth, thank you for your outstanding review. I particularly like the way you got rid of the duplicated logic by introducing the create parameter in _find, though I think I might make it keyword-only.
• Winston Ewert: It only makes sense to subclass dict if you are making some kind of special dict. Here you are creating a class which uses a dict, and I think thats much more naturally represented using composition rather then inheritance. I'll admit, its a bit a purest thing here.

Answer 2

A note on terminology: a mapping is, according to Python glossary,

A container object that supports arbitrary key lookups and implements the methods specified in the Mapping or MutableMapping abstract base classes. Examples include dict, collections.defaultdict, collections.OrderedDict and collections.Counter.

Your from_mapping assumes a different type of argument.

• Adam: Thanks, that's a very good point. In fact, I was confused by the dict constructor documentation but that's because I didn't read it carefully enough. Any suggestions for a better name? I'm thinking of key_value_pairs...
• Winston Ewert: @codesparkle, I'd go from from_pairs myself.
• Janne Karila: @codesparkle I'd follow Gareth's suggestion and have the constructor call update.

Answer 3

__getitem__ returns None if a key cannot be found, to follow python conventions it should raise KeyError.

Simliar with __delitem__, I'd expected a KeyError if deleting a non-existent element.

I wonder if it would be better to have your _Entry object support a .items() method which could then be implemented recursively.

Given that it act a lot like a mapping, I'd expect the update method to act like dict.update but it doesn't.

 if entry.value != new_value:
 entry.value = new_value
 self._length += 1

Its not clear to me that this is the correct logic for incrementing the length.

• Adam: Thank you for your excellent points. You're right about length: it is indeed a bug.