blog/content/posts/2024-06-30-trie/index.md

title	date	draft	description	tags	categories	series	favorite	disable_feed
Trie	2024-06-30T11:07:49+01:00	false	A cool map	algorithms data structures python	programming	Cool algorithms	false	false

This time, let's talk about the Trie, which is a tree-based mapping structure most often used for string keys.

What does it do?

A Trie can be used to map a set of string keys to their corresponding values, without the need for a hash function. This also means you won't suffer from hash collisions, though the tree-based structure will probably translate to slower performance than a good hash table.

A Trie is especially useful to represent a dictionary of words in the case of spell correction, as it can easily be used to fuzzy match words under a given edit distance (think Levenshtein distance)

Implementation

This implementation will be in Python for exposition purposes, even though it already has a built-in dict.

Representation

Creating a new Trie is easy: the root node starts off empty and without any mapped values.

class Trie[T]:
    _children: dict[str, Trie[T]]
    _value: T | None

    def __init__(self):
        # Each letter is mapped to a Trie
        self._children = defaultdict(Trie)
        # If we match a full string, we store the mapped value
        self._value = None

We're using a defaultdict for the children for ease of implementation in this post. In reality, I would encourage you exit early when you can't match a given character.

The string key will be implicit by the position of a node in the tree: the empty string at the root, one-character strings as its direct children, etc...

Search

An exact match look-up is easily done: we go down the tree until we've exhausted the key. At that point we've either found a mapped value or not.

def get(self, key: str) -> T | None:
    # Have we matched the full key?
    if not key:
        # Store the `T` if mapped, `None` otherwise
        return self._value
    # Otherwise, recurse on the child corresponding to the first letter
    return self._children[key[0]].get(key[1:])

Insertion

Adding a new value to the Trie is similar to a key lookup, only this time we store the new value instead of returning it.

def insert(self, key: str, value: T) -> bool:
    # Have we matched the full key?
    if not key:
        # Check whether we're overwriting a previous mapping
        was_mapped = self._value is None
        # Store the corresponding value
        self._value = value
        # Return whether we've performed an overwrite
        return was_mapped
      # Otherwise, recurse on the child corresponding to the first letter
      return self._children[key[0]].insert(key[1:], value)

Removal

Removal should also look familiar.

def remove(self, key: str) -> bool:
    # Have we matched the full key?
    if not key:
        was_mapped = self._value is None
        # Remove the value
        self._value = None
        # Return whether it was mapped
        return was_mapped
    # Otherwise, recurse on the child corresponding to the first letter
    return self._children[key[0]].remove(key[1:])

Fuzzy matching

Fuzzily matching a given word is where the real difficulty is: the key is to realize we can use the prefix-tree nature of a Trie to avoid doing wasteful work.

By leveraging the prefix visit order of the tree, we can build an iterative Levenshtein distance matrix, in much the same way one would do so in its Dynamic Programming implementation (see the Wagner-Fisher algorithm).

class FuzzyResult[T](NamedTuple):
    distance: int
    key: str
    value: T


def get_fuzzy(self, key: str, max_distance: int = 0) -> Iterator[FuzzyResult[T]]:
    def helper(
        current_word: str,
        node: Trie[T],
        previous_row: list[int],
    ) -> Iterator[tuple[int, T]]:
        # Iterative Levenshtein
        current_row = [previous_row[0] + 1]
        current_char = current_word[-1]
        for column, key_char in enumerate(key, start=1):
            insertion = current_row[column - 1] + 1
            deletion = previous_row[column] + 1
            replacement = previous_row[column - 1] + (key_char != current_char)
            current_row.append(min(insertion, deletion, replacement))

        # If we are under the max distance, match this node
        if (distance := current_row[-1]) <= max_distance and node._value != None:
            # Only if it has a value of course
            yield FuzzyResult(distance, current_word, node._value)

        # If we can potentially still match children, recurse
        if min(current_row) <= max_distance:
            for c, child in node._children.items():
                yield from helper(current_word + c, child, current_row)

    # Build the first row -- the edit distance from the empty string
    row = list(range(len(key) + 1))

    # Base case for the empty string
    if (distance := row[-1]) <= max_distance and self._value != None:
        yield FuzzyResult(distance, "", self._value)
    for c, child in self._children.items():
        yield from helper(c, child, row)