Trie Tree

We have a bunch of key-value data. When someone types a substring of a key, we want to easily retrieve all key-value pairs whose keys contain that substring. A Trie tree is a nice solution for this.

As the name suggests, a Trie tree is a trie. The general idea is to build a trie like the one shown below:

It is used to store strings such as to, tea, ted, ten, A, in, inn, and so on.

A query starts from the root, iterates over each character of the query string, and transitions state based on that character, until the transition fails or the query string is exhausted. At that point, everything in the subtree rooted at that node contains the query string. In this sense, a query is a finite state machine: reach the terminal state and then take the subtree.

When you type t, you jump from the root to the t node; then when you type e, you jump to the e node.

If your query string is te, you stop at the te node, and all its descendants (tea, ted, ten) are the results.

But, as many sources point out, it is a prefix tree. If we query ea, we won’t get the desired tea at this point. To support arbitrary substrings, we need an extra index: instead of starting from the root, we start from the index, query, and jump to all nodes that begin with the first character. That makes this easy to implement.

There are many implementation choices for the tree, depending on the specific use case, so implementations naturally vary. In my case, I assume the same key may have different values, and entering a substring returns all corresponding values. In practice, my key-value pairs number around 350k, so I lean toward a conservative time-space trade-off.

I posted one of my implementations to a gist; the link is here. A brief introduction follows.

First, the node definition is as follows:

case class TrieNode[V](
    base: Option[Char], // current node
    check: Option[TrieNode[V]], // previous node , None for the first node
    next: TrieMap[Char, TrieNode[V]]) {
  var v: MutableSet[V] = MutableSet[V]()
}

This is a simple triple structure. v holds the associated values; the other fields are needed to build the tree: base, the current node’s state value; check, the parent node’s address; and next, the map of child nodes.

By the way, I originally wanted to see if Scala had a built-in trie tree, but instead found TrieMap. I gave it a try and it felt pretty handy.

Then the code for adding content:

  def put(k: String, v: V) = insertMonitor.synchronized {
    var p = root
    for (c <- k) {
      if (!p.next.contains(c)) {
        val t = (c -> new TrieNode[V](Some(c), Some(p), TrieMap[Char, TrieNode[V]]()))
        index.getOrElseUpdate(c, ArrayBuffer[TrieNode[V]]()) += t._2
        p.next += t
      }
      p = p.next.get(c).get
    }
    p.v += v
  }

root is the root node, defined as TrieNode. The index is defined as TrieMap[Char, ArrayBuffer[TrieNode[V]]]. The idea is to walk the key through the state machine starting from the root and leave the value at the terminal state.

Queries are also convenient:

  def find(k: String, limit: Option[Int] = None): List[V] = {
    if (k.length() > 0) {
      val tk = k.splitAt(1)
      var pSet = index.get(tk._1.charAt(0)) match {
        case Some(ps) => ps
        case None => ArrayBuffer[TrieNode[V]]()
      }
      val rpSet = pSet.par.flatMap { p =>
        var rp = Option(p)
        def checkString {
          for (c <- tk._2) {
            rp = rp match {
              case Some(gp) => gp.next.get(c)
              case None => return
            }
          }
        }
        checkString
        rp
      }
      //ignored: get results
      List[V]()
    } else {
      List[V]()
    }
  }

In short, it looks up the list of states reachable from the index by the first character, then transitions through each state to collect the values under the subtree. (Details are omitted here for brevity; see the gist.)

Trie trees have many issues; first, memory usage is relatively high. To avoid storing values redundantly, I only place values at the final state. Therefore, when the result set is large—for example, querying a single s may return tens of thousands of results—it takes considerable time (about one to two seconds on my 300k dataset), which is a drawback.

All in all, it’s a choice between spending space or spending time.