LeetCode Meditations — Chapter 10: Tries

The trie data structure gets its name from the word retrieval — and it's usually pronounced as "try," so that we don't get confused with another familiar and friendly data structure, "tree."

However, a trie is still a tree (or tree-like) data structure whose nodes usually store individual letters. So, by traversing the nodes in a trie, we can retrieve strings.

Tries are useful for applications such as autocompletion and spellchecking — and the larger our trie is, the less work we have to do for inserting a new value.

---

An important note before we start: using arrays is not very memory-efficient, and we'll see another way of creating tries in the next article for Implement Trie (Prefix Tree). For now, we'll stick to the array implementation.

First, let's see what a trie looks like:

In this trie, we can retrieve the strings "sea" and "see" — but not "sew" for example.

There is a lot going on, but we can try to understand it piece by piece.

Let's look at a trie node.

We'll create a TrieNode class that has children, which is an array of length 26 (so that each index corresponds to a letter in the English alphabet), and a flag variable isEndOfWord to indicate whether that node represents the last character of a word:

class TrieNode {
  constructor() {
    this.children = Array.from({ length: 26 }, () => null);
    this.isEndOfWord = false;
  }
}

We're initializing children with null values. As we add a character to our trie, the index that corresponds to that character will be filled.

Note
We're not storing the actual character itself in this implementation, it's implicit in the usage of indices.

In a trie, we start with an empty root node.

class Trie {
  constructor() {
    this.root = new TrieNode();
  }
  ...
}

To insert a word, we're going to loop through each character, and initialize a new TrieNode to the corresponding index.

insert(word) {
  let currentNode = this.root;
  for (const char of word) {
    let idx = char.charCodeAt(0) - 'a'.charCodeAt(0);
    if (currentNode.children[idx] === null) {
      currentNode.children[idx] = new TrieNode();
    }
    currentNode = currentNode.children[idx];
  }

  currentNode.isEndOfWord = true;
}

Once we reach the node that indicates the last character of the word we inserted, we also mark the isEndOfWord variable as true.

Note
word is going to be lowercase in these examples, otherwise, we have to convert it, such as: word = word.toLowerCase();

For searching a word's existence in the trie, we'll do a similar thing. We'll look at the nodes for each character, and if we reach the last one that has isEndOfWord marked as true, that means we've found the word:

search(word) {
  let currentNode = this.root;
  for (const char of word) {
    let idx = char.charCodeAt(0) - 'a'.charCodeAt(0);
    if (currentNode.children[idx] === null) {
      return false;
    }      
    currentNode = currentNode.children[idx];
  }

  return currentNode.isEndOfWord;
}

Note
If we find the word we're looking for, then it's called a search hit; otherwise, we have a search miss and the word doesn't exist in our trie.

Removing a word is a bit more challenging. Let's say that we want to remove the word "see." But, there is also another word "sea," with the same prefix ('s' and 'e'). So, we should remove only the nodes that we're allowed to.

For this reason, we'll define a recursive function.
Once we reach the last character of the word we want to remove, we'll back up and remove the characters we can remove:

const removeRecursively = (node, word, depth) => {
  if (node === null) {
    return null;
  }

  if (depth === word.length) {
    if (node.isEndOfWord) {
      node.isEndOfWord = false;
    }
    if (node.children.every(child => child === null)) {
      node = null;
    }

    return node;
  }

  let idx = word[depth].charCodeAt(0) - 'a'.charCodeAt(0);
  node.children[idx] = removeRecursively(node.children[idx], word, depth + 1);

  if (node.children.every(child => child === null) && !node.isEndOfWord) {
    node = null;
  }

  return node;
}

depth indicates the index of the word, or the depth of the trie we reach.

Once depth is equal to the word's length (one past the last character), we check if it's the end of the word, if that's the case, we'll mark it as false now, because that word won't exist from here on. Then, we can only mark the node as null if it doesn't have any children (in other words, if all of them are null). We'll apply this logic to each child node recursively until the word is removed as far as it can be removed.

Here is the final example implementation of a trie:

class TrieNode {
  constructor() {
    this.children = Array.from({ length: 26 }, () => null);
    this.isEndOfWord = false;
  }
}

class Trie {
  constructor() {
    this.root = new TrieNode();
  }

  insert(word) {
    let currentNode = this.root;
    for (const char of word) {
      let idx = char.charCodeAt(0) - 'a'.charCodeAt(0);
      if (currentNode.children[idx] === null) {
        currentNode.children[idx] = new TrieNode();
      }
      currentNode = currentNode.children[idx];
    }

    currentNode.isEndOfWord = true;
  }

  search(word) {
    let currentNode = this.root;
    for (const char of word) {
      let idx = char.charCodeAt(0) - 'a'.charCodeAt(0);
      if (currentNode.children[idx] === null) {
        return false;
      }      
      currentNode = currentNode.children[idx];
    }

    return currentNode.isEndOfWord;
  }

  remove(word) {
    const removeRecursively = (node, word, depth) => {
      if (node === null) {
        return null;
      }

      if (depth === word.length) {
        if (node.isEndOfWord) {
          node.isEndOfWord = false;
        }
        if (node.children.every(child => child === null)) {
          node = null;
        }

        return node;
      }

      let idx = word[depth].charCodeAt(0) - 'a'.charCodeAt(0);
      node.children[idx] = removeRecursively(node.children[idx], word, depth + 1);

      if (node.children.every(child => child === null) && !node.isEndOfWord) {
        node = null;
      }

      return node;
    }

    removeRecursively(this.root, word, 0);
  }
}

let t = new Trie();

t.insert('sea');
t.insert('see');

console.log(t.search('sea')); // true
console.log(t.search('see')); // true

console.log(t.search('hey')); // false
console.log(t.search('sew')); // false

t.remove('see');

console.log(t.search('see')); // false 
console.log(t.search('sea')); // true

---

Time and space complexity

The time complexity of creating a trie is going to be $O(m * n)$ where $m$ is the longest word and $n$ is the total number of words.
Inserting, searching, and deleting a word is $O(a * n)$ where $a$ is the length of the word and $n$ is the total number of words.

When it comes to space complexity, in the worst case, each node can have children for all the characters in the alphabet we're representing. But, the size of the alphabet is constant, so the growth of storage needs will be proportionate to the number of nodes we have, which is $O(n)$ where $n$ is the number of nodes.

---

We have already done most of the work for the next problem, but next time we'll be slightly more efficient. Until then, happy coding.