Implementing Tries in Python

In the article, we introduced the prefix tree data structure (better known as a trie), its properties, and some of its applications. In this lesson, we will see how we can implement a trie.

Recall that a trie is an efficient way to store a collection of strings by storing them in a tree-like data structure.

We will create a Python object that will represent the trie data structure. This object will contain functions to store a word, search for a word, and count the number of words that contain a particular prefix. 

We will also create a Python object that represents a node in the trie.

We will construct the Python class to the right and define all necessary functions. Then we will see an application of the trie as a spell checker.

Introduction

Great job on making it this far in the lesson on Tries! This is a very powerful data structure with many applications in software engineering, computer science, and coding interviews!

Let’s summarize everything we went over:

* A trie is a data structure that efficiently stores strings based on their prefix.
* A trie is a tree-like data structure in which each node represents a character.
* Each node’s child contains other nodes that represent subsequent characters in a string.
* The `Trie` data structure we created supported the following functions:
  * Adding a word
  * Searching for a word
  * Counting a prefix
* Each of these functions traverses the trie starting from the root and goes through the nodes corresponding to the characters in a string (prefix or word).
* Relevant data is stored in the trie’s nodes.
* A trie can be used as the base data structure for a spell checker.

Good work!


Review

In this exercise, we will work on the `count_prefix()` method in our `Trie` class. As its name suggests, this function will count the number of words with a prefix up to some character. We track this count in the `freq` variable of a `TrieNode`.

Consider this example of a trie that contains the following words:
* ADD
* ADDITION
* ADDENDUM
* ADVERTISEMENT
* ADVANCED
* APPLE

The `freq` variable in the node for character 'A' will be set to 6 since all six words have a prefix 'A'. Node 'D' will have `freq` set to 5 because five words have the prefix "AD". The second node 'D' will be set to 3. Node 'V' will have `freq` set to 2 because only two words have the prefix `ADV`. The counter in Node 'P' will be set 1 because of APPLE.

The algorithm for this process is the same algorithm as that for `search_string()`. The only difference is that `count_prefix()` will return the value of the last character's `freq` variable instead of returning the value of the `end_of_key` variable.


Searching for a Prefix in a Trie

In this exercise, we will describe the algorithm that adds a word to a trie. We will then implement this algorithm in the `add_string()` function of the `Trie` class. We will be implementing this algorithm iteratively.

Let's go over the algorithm to add a string to a trie step-by-step:
* Begin by creating a pointer that initially points to the root node of the `Trie`.
* Iterate through the characters of the `string` parameter using a loop.
* Take the first character and call the `root` node's `add_char()` method passing in the character as a parameter.
* The `add_char()` method processes the character according to its logic.
* The `add_char()` method then returns the node associated with the current character.
* The pointer will now point to the node returned from `add_char()`.
* Before the loop continues, the returned node associated with the character has a `freq` variable that is keeping track of the number of prefixes up to this character. `freq` is now incremented.
* The loop repeats this process for the remaining characters.
* The very last node, representing the last character in the string, will have its `end_of_key` set to `True` to indicate that this is the end of a string.

To review this, let's take a look at an example. We wish to add the word "CODE" to the trie:
1. We call `add_word(CODE)`.
2. Inside `add_word()`, we initialize a pointer to the `root` of the trie.
3. A loop is created that will iterate through "CODE".
4. We will call `root.add_char('C')`.
5. The `root.add_char()` function will create an entry in the `root` node’s `nodes` dictionary if one does not already exist.
6. `root.add_char()` will return the node that is associated with the character 'C'.
7. We set the pointer to this node.
8. We increment the `freq` variable by one to represent that this node (this character) is part of `freq` amount of prefixes.
9. The loop proceeds on to the next letter 'O' and repeats steps 2-9 for the remaining characters.
10. In the end, the node representing the character 'E' will have its `end_of_key` variable set to `True`. 

Let's implement this algorithm in the `add_string()` method!


Adding a String to a Trie

Let's define the `TrieNode` class we will use to implement a trie.

To represent subsequent characters in a string, a `TrieNode` object contains a dictionary that maps a character to an instance of `TrieNode`. This works somewhat like a linked list in which you have a node that contains information and a pointer to another node.

We represent the `TrieNode` as the following clas:

```py
class TrieNode:

  def __init__(self):
    self.nodes = {}
```

The `self.nodes` field is a dictionary that contains the mapping of each character and its subsequent node mapping.

`TrieNode` also contains a function used to append a character:

```py
class TrieNode:

  def __init__(self):
    self.nodes = {}

  def add_char(c: chr) -> TrieNode:
    if c not in self.nodes:
      self.nodes[c] = TrieNode()
    return self.nodes[c]
```

We call this method from the `Trie` class. It returns a `TrieNode` because a loop in the `add_word()` method will be responsible for iterating through the nodes starting at the `root` and placing the characters where they need to be. There are two possible scenarios: 
* The current `TrieNode` already contains an entry for character `c` because it has been entered earlier as the prefix to another word. In this case, `add_char()` retrieves the `TrieNode` mapped to `c` and returns it. 
* The current prefix has not been entered. Therefore, `TrieNode` does not contain a mapping for the character `c`. In this case, `add_char()` creates a mapping and returns it. `add_word()` will continue to load the word into the trie.

We can augment the `TrieNode` class to hold other relevant information about a word.


The Trie Node

Now that you have finished constructing your trie data structure, in this exercise, you will get a chance to see it in action.

In the following checkpoints, you will add strings, search for strings, and count prefixes. 

Using a Trie

In this exercise, we will search for a string in a trie. A string is loaded into a trie for one of two reasons:
* The trie represents a dictionary of words.
* The string contains some relevant information.  

The function `search_string()` in our `Trie` class will search the trie for the string and return whether it is there or not. We can modify this function to return relevant data associated with the string, but for our purposes, we will only determine whether the string is in the trie or not.

Here are the steps of the algorithm to search for a string:
* Initialize a pointer to the root node in the trie.
* Create a loop that will loop through all the characters of the string passed into `search_string()`.
* Access the current node's `nodes` dictionary and attempt to retrieve the node associated with the current character.
    * If the node exists, set the pointer to this node.
    * If the node does not exist, that means this string is not there; therefore, the loop should terminate immediately, and `search_string()` should return `False`.
* Assuming the word exists, the loop will repeat this process for all remaining characters.
* At the end of the loop, the pointer points to the node associated with the last character in the string; therefore, the function should return the value associated with `end_of_key`.

Please note that we return the value of `end_of_key` and not `True` because we want to know if this string is a full entry. A string may exist in the trie only because it is the prefix to a longer string. In that case, we should return `False` because the string has not been entered as its own word.

Let's take a look at an example. We have a trie that contains the word CODECADEMY.

We will search this trie for the word CODE (this word has *not* been entered into the trie):

1. Initialize a pointer to the root node.
2. A loop will loop through all characters of CODE.
3. We will access the root node's `nodes` dictionary to retrieve the node associated with the letter 'C'.
4. We will set the pointer to that node.
5. Repeat steps 2 through 4 for the remaining characters.
6. This loop will not terminate early because the string CODE exists in the trie as a prefix of CODECADEMY.
7. After the loop terminates, the pointer will point to the node associated with the letter 'E'.
8. Return the value of the `end_of_key` variable of the node associated with the letter 'E'. This will be `False` because CODE does not exist in the trie as its own word. 

If you search for the word CODECADEMY, `search_string()` will return `True`.

If you search for the word COCOA, `search_string()` will terminate early and return `False` because there is no node associated with the second 'C'.

If there is relevant information associated with the word, it will likely be stored as a field in one of the `TrieNodes`. If you want that, have `search_string()` return it!
